1# coding: utf-8
2# Copyright (c) Pymatgen Development Team.
3# Distributed under the terms of the MIT License.
4
5
6"""
7This module contains some math utils that are used in the chemenv package.
8"""
9
10__author__ = "David Waroquiers"
11__copyright__ = "Copyright 2012, The Materials Project"
12__credits__ = "Geoffroy Hautier"
13__version__ = "2.0"
14__maintainer__ = "David Waroquiers"
15__email__ = "david.waroquiers@gmail.com"
16__date__ = "Feb 20, 2016"
17
18from functools import reduce
19from math import sqrt
20
21import numpy as np
22from scipy.special import erf
23
24##############################################################
25# cartesian product of lists ##################################
26##############################################################
27
28
29def _append_es2sequences(sequences, es):
30    result = []
31    if not sequences:
32        for e in es:
33            result.append([e])
34    else:
35        for e in es:
36            result += [seq + [e] for seq in sequences]
37    return result
38
39
40def _cartesian_product(lists):
41    """
42    given a list of lists,
43    returns all the possible combinations taking one element from each list
44    The list does not have to be of equal length
45    """
46    return reduce(_append_es2sequences, lists, [])
47
48
49def prime_factors(n):
50    """Lists prime factors of a given natural integer, from greatest to smallest
51    :param n: Natural integer
52    :rtype : list of all prime factors of the given natural n
53    """
54    i = 2
55    while i <= sqrt(n):
56        if n % i == 0:
57            l = prime_factors(n / i)
58            l.append(i)
59            return l
60        i += 1
61    return [n]  # n is prime
62
63
64def _factor_generator(n):
65    """
66    From a given natural integer, returns the prime factors and their multiplicity
67    :param n: Natural integer
68    :return:
69    """
70    p = prime_factors(n)
71    factors = {}
72    for p1 in p:
73        try:
74            factors[p1] += 1
75        except KeyError:
76            factors[p1] = 1
77    return factors
78
79
80def divisors(n):
81    """
82    From a given natural integer, returns the list of divisors in ascending order
83    :param n: Natural integer
84    :return: List of divisors of n in ascending order
85    """
86    factors = _factor_generator(n)
87    _divisors = []
88    listexponents = [[k ** x for x in range(0, factors[k] + 1)] for k in list(factors.keys())]
89    listfactors = _cartesian_product(listexponents)
90    for f in listfactors:
91        _divisors.append(reduce(lambda x, y: x * y, f, 1))
92    _divisors.sort()
93    return _divisors
94
95
96def get_center_of_arc(p1, p2, radius):
97    """
98    :param p1:
99    :param p2:
100    :param radius:
101    :return:
102    """
103    dx = p2[0] - p1[0]
104    dy = p2[1] - p1[1]
105    dd = np.sqrt(dx * dx + dy * dy)
106    radical = np.power((radius / dd), 2) - 0.25
107    if radical < 0:
108        raise ValueError("Impossible to find center of arc because the arc is ill-defined")
109    tt = np.sqrt(radical)
110    if radius > 0:
111        tt = -tt
112    return (p1[0] + p2[0]) / 2 - tt * dy, (p1[1] + p2[1]) / 2 + tt * dx
113
114
115def get_linearly_independent_vectors(vectors_list):
116    """
117    :param vectors_list:
118    :return:
119    """
120    independent_vectors_list = []
121    for vector in vectors_list:
122        if np.any(vector != 0):
123            if len(independent_vectors_list) == 0:
124                independent_vectors_list.append(np.array(vector))
125            elif len(independent_vectors_list) == 1:
126                rank = np.linalg.matrix_rank(np.array([independent_vectors_list[0], vector, [0, 0, 0]]))
127                if rank == 2:
128                    independent_vectors_list.append(np.array(vector))
129            elif len(independent_vectors_list) == 2:
130                mm = np.array([independent_vectors_list[0], independent_vectors_list[1], vector])
131                if np.linalg.det(mm) != 0:
132                    independent_vectors_list.append(np.array(vector))
133        if len(independent_vectors_list) == 3:
134            break
135    return independent_vectors_list
136
137
138def scale_and_clamp(xx, edge0, edge1, clamp0, clamp1):
139    """
140    :param xx:
141    :param edge0:
142    :param edge1:
143    :param clamp0:
144    :param clamp1:
145    :return:
146    """
147    return np.clip((xx - edge0) / (edge1 - edge0), clamp0, clamp1)
148
149
150# Step function based on the cumulative distribution function of the normal law
151def normal_cdf_step(xx, mean, scale):
152    """
153    :param xx:
154    :param mean:
155    :param scale:
156    :return:
157    """
158    return 0.5 * (1.0 + erf((xx - mean) / (np.sqrt(2.0) * scale)))
159
160
161# SMOOTH STEP FUNCTIONS
162# Set of smooth step functions that allow to smoothly go from y = 0.0 (1.0) to y = 1.0 (0.0) by changing x
163# from 0.0 to 1.0 respectively when inverse is False (True).
164# (except if edges is given in which case a the values are first scaled and clamped to the interval given by edges)
165# The derivative at x = 0.0 and x = 1.0 have to be 0.0
166
167
168def smoothstep(xx, edges=None, inverse=False):
169    """
170    :param xx:
171    :param edges:
172    :param inverse:
173    :return:
174    """
175    if edges is None:
176        xx_clipped = np.clip(xx, 0.0, 1.0)
177        if inverse:
178            return 1.0 - xx_clipped * xx_clipped * (3.0 - 2.0 * xx_clipped)
179        return xx_clipped * xx_clipped * (3.0 - 2.0 * xx_clipped)
180    xx_scaled_and_clamped = scale_and_clamp(xx, edges[0], edges[1], 0.0, 1.0)
181    return smoothstep(xx_scaled_and_clamped, inverse=inverse)
182
183
184def smootherstep(xx, edges=None, inverse=False):
185    """
186    :param xx:
187    :param edges:
188    :param inverse:
189    :return:
190    """
191    if edges is None:
192        xx_clipped = np.clip(xx, 0.0, 1.0)
193        if inverse:
194            return 1.0 - xx_clipped * xx_clipped * xx_clipped * (xx_clipped * (xx_clipped * 6 - 15) + 10)
195        return xx_clipped * xx_clipped * xx_clipped * (xx_clipped * (xx_clipped * 6 - 15) + 10)
196    xx_scaled_and_clamped = scale_and_clamp(xx, edges[0], edges[1], 0.0, 1.0)
197    return smootherstep(xx_scaled_and_clamped, inverse=inverse)
198
199
200def cosinus_step(xx, edges=None, inverse=False):
201    """
202    :param xx:
203    :param edges:
204    :param inverse:
205    :return:
206    """
207    if edges is None:
208        xx_clipped = np.clip(xx, 0.0, 1.0)
209        if inverse:
210            return (np.cos(xx_clipped * np.pi) + 1.0) / 2.0
211        return 1.0 - (np.cos(xx_clipped * np.pi) + 1.0) / 2.0
212    xx_scaled_and_clamped = scale_and_clamp(xx, edges[0], edges[1], 0.0, 1.0)
213    return cosinus_step(xx_scaled_and_clamped, inverse=inverse)
214
215
216def power3_step(xx, edges=None, inverse=False):
217    """
218    :param xx:
219    :param edges:
220    :param inverse:
221    :return:
222    """
223    return smoothstep(xx, edges=edges, inverse=inverse)
224
225
226def powern_parts_step(xx, edges=None, inverse=False, nn=2):
227    """
228    :param xx:
229    :param edges:
230    :param inverse:
231    :param nn:
232    :return:
233    """
234    if edges is None:
235        aa = np.power(0.5, 1.0 - nn)
236        xx_clipped = np.clip(xx, 0.0, 1.0)
237        if np.mod(nn, 2) == 0:
238            if inverse:
239                return 1.0 - np.where(
240                    xx_clipped < 0.5,
241                    aa * np.power(xx_clipped, nn),
242                    1.0 - aa * np.power(xx_clipped - 1.0, nn),
243                )
244            return np.where(
245                xx_clipped < 0.5,
246                aa * np.power(xx_clipped, nn),
247                1.0 - aa * np.power(xx_clipped - 1.0, nn),
248            )
249        if inverse:
250            return 1.0 - np.where(
251                xx_clipped < 0.5,
252                aa * np.power(xx_clipped, nn),
253                1.0 + aa * np.power(xx_clipped - 1.0, nn),
254            )
255        return np.where(
256            xx_clipped < 0.5,
257            aa * np.power(xx_clipped, nn),
258            1.0 + aa * np.power(xx_clipped - 1.0, nn),
259        )
260    xx_scaled_and_clamped = scale_and_clamp(xx, edges[0], edges[1], 0.0, 1.0)
261    return powern_parts_step(xx_scaled_and_clamped, inverse=inverse, nn=nn)
262
263
264# FINITE DECREASING FUNCTIONS
265# Set of decreasing functions that allow to smoothly go from y = 1.0 to y = 0.0 by changing x from 0.0 to 1.0
266# The derivative at x = 1.0 has to be 0.0
267
268
269def powern_decreasing(xx, edges=None, nn=2):
270    """
271    :param xx:
272    :param edges:
273    :param nn:
274    :return:
275    """
276    if edges is None:
277        aa = 1.0 / np.power(-1.0, nn)
278        return aa * np.power(xx - 1.0, nn)
279    xx_scaled_and_clamped = scale_and_clamp(xx, edges[0], edges[1], 0.0, 1.0)
280    return powern_decreasing(xx_scaled_and_clamped, nn=nn)
281
282
283def power2_decreasing_exp(xx, edges=None, alpha=1.0):
284    """
285    :param xx:
286    :param edges:
287    :param alpha:
288    :return:
289    """
290    if edges is None:
291        aa = 1.0 / np.power(-1.0, 2)
292        return aa * np.power(xx - 1.0, 2) * np.exp(-alpha * xx)
293    xx_scaled_and_clamped = scale_and_clamp(xx, edges[0], edges[1], 0.0, 1.0)
294    return power2_decreasing_exp(xx_scaled_and_clamped, alpha=alpha)
295
296
297# INFINITE TO FINITE DECREASING FUNCTIONS
298# Set of decreasing functions that allow to smoothly go from y = + Inf to y = 0.0 by changing x from 0.0 to 1.0
299# The derivative at x = 1.0 has to be 0.0
300
301
302def power2_tangent_decreasing(xx, edges=None, prefactor=None):
303    """
304    :param xx:
305    :param edges:
306    :param prefactor:
307    :return:
308    """
309    if edges is None:
310        if prefactor is None:
311            aa = 1.0 / np.power(-1.0, 2)
312        else:
313            aa = prefactor
314        return -aa * np.power(xx - 1.0, 2) * np.tan((xx - 1.0) * np.pi / 2.0)  # pylint: disable=E1130
315
316    xx_scaled_and_clamped = scale_and_clamp(xx, edges[0], edges[1], 0.0, 1.0)
317    return power2_tangent_decreasing(xx_scaled_and_clamped, prefactor=prefactor)
318
319
320def power2_inverse_decreasing(xx, edges=None, prefactor=None):
321    """
322    :param xx:
323    :param edges:
324    :param prefactor:
325    :return:
326    """
327    if edges is None:
328        if prefactor is None:
329            aa = 1.0 / np.power(-1.0, 2)
330        else:
331            aa = prefactor
332        return np.where(np.isclose(xx, 0.0), aa * float("inf"), aa * np.power(xx - 1.0, 2) / xx)
333        # return aa * np.power(xx-1.0, 2) / xx if xx != 0 else aa * float("inf")
334    xx_scaled_and_clamped = scale_and_clamp(xx, edges[0], edges[1], 0.0, 1.0)
335    return power2_inverse_decreasing(xx_scaled_and_clamped, prefactor=prefactor)
336
337
338def power2_inverse_power2_decreasing(xx, edges=None, prefactor=None):
339    """
340    :param xx:
341    :param edges:
342    :param prefactor:
343    :return:
344    """
345    if edges is None:
346        if prefactor is None:
347            aa = 1.0 / np.power(-1.0, 2)
348        else:
349            aa = prefactor
350        return np.where(
351            np.isclose(xx, 0.0),
352            aa * float("inf"),
353            aa * np.power(xx - 1.0, 2) / xx ** 2.0,
354        )
355    xx_scaled_and_clamped = scale_and_clamp(xx, edges[0], edges[1], 0.0, 1.0)
356    return power2_inverse_power2_decreasing(xx_scaled_and_clamped, prefactor=prefactor)
357
358
359def power2_inverse_powern_decreasing(xx, edges=None, prefactor=None, powern=2.0):
360    """
361    :param xx:
362    :param edges:
363    :param prefactor:
364    :param powern:
365    :return:
366    """
367    if edges is None:
368        if prefactor is None:
369            aa = 1.0 / np.power(-1.0, 2)
370        else:
371            aa = prefactor
372        return aa * np.power(xx - 1.0, 2) / xx ** powern
373
374    xx_scaled_and_clamped = scale_and_clamp(xx, edges[0], edges[1], 0.0, 1.0)
375    return power2_inverse_powern_decreasing(xx_scaled_and_clamped, prefactor=prefactor, powern=powern)
376