diff --git a/lib/python3.12/site-packages/sympy/physics/mechanics/functions.py b/lib/python3.12/site-packages/sympy/physics/mechanics/functions.py new file mode 100644 index 0000000000000000000000000000000000000000..42abe2b7fe608b4602cdab518f209b446b2dbe03 --- /dev/null +++ b/lib/python3.12/site-packages/sympy/physics/mechanics/functions.py @@ -0,0 +1,735 @@ +from sympy.utilities import dict_merge +from sympy.utilities.iterables import iterable +from sympy.physics.vector import (Dyadic, Vector, ReferenceFrame, + Point, dynamicsymbols) +from sympy.physics.vector.printing import (vprint, vsprint, vpprint, vlatex, + init_vprinting) +from sympy.physics.mechanics.particle import Particle +from sympy.physics.mechanics.rigidbody import RigidBody +from sympy.simplify.simplify import simplify +from sympy import Matrix, Mul, Derivative, sin, cos, tan, S +from sympy.core.function import AppliedUndef +from sympy.physics.mechanics.inertia import (inertia as _inertia, + inertia_of_point_mass as _inertia_of_point_mass) +from sympy.utilities.exceptions import sympy_deprecation_warning + +__all__ = ['linear_momentum', + 'angular_momentum', + 'kinetic_energy', + 'potential_energy', + 'Lagrangian', + 'mechanics_printing', + 'mprint', + 'msprint', + 'mpprint', + 'mlatex', + 'msubs', + 'find_dynamicsymbols'] + +# These are functions that we've moved and renamed during extracting the +# basic vector calculus code from the mechanics packages. + +mprint = vprint +msprint = vsprint +mpprint = vpprint +mlatex = vlatex + + +def mechanics_printing(**kwargs): + """ + Initializes time derivative printing for all SymPy objects in + mechanics module. + """ + + init_vprinting(**kwargs) + +mechanics_printing.__doc__ = init_vprinting.__doc__ + + +def inertia(frame, ixx, iyy, izz, ixy=0, iyz=0, izx=0): + sympy_deprecation_warning( + """ + The inertia function has been moved. + Import it from "sympy.physics.mechanics". + """, + deprecated_since_version="1.13", + active_deprecations_target="moved-mechanics-functions" + ) + return _inertia(frame, ixx, iyy, izz, ixy, iyz, izx) + + +def inertia_of_point_mass(mass, pos_vec, frame): + sympy_deprecation_warning( + """ + The inertia_of_point_mass function has been moved. + Import it from "sympy.physics.mechanics". + """, + deprecated_since_version="1.13", + active_deprecations_target="moved-mechanics-functions" + ) + return _inertia_of_point_mass(mass, pos_vec, frame) + + +def linear_momentum(frame, *body): + """Linear momentum of the system. + + Explanation + =========== + + This function returns the linear momentum of a system of Particle's and/or + RigidBody's. The linear momentum of a system is equal to the vector sum of + the linear momentum of its constituents. Consider a system, S, comprised of + a rigid body, A, and a particle, P. The linear momentum of the system, L, + is equal to the vector sum of the linear momentum of the particle, L1, and + the linear momentum of the rigid body, L2, i.e. + + L = L1 + L2 + + Parameters + ========== + + frame : ReferenceFrame + The frame in which linear momentum is desired. + body1, body2, body3... : Particle and/or RigidBody + The body (or bodies) whose linear momentum is required. + + Examples + ======== + + >>> from sympy.physics.mechanics import Point, Particle, ReferenceFrame + >>> from sympy.physics.mechanics import RigidBody, outer, linear_momentum + >>> N = ReferenceFrame('N') + >>> P = Point('P') + >>> P.set_vel(N, 10 * N.x) + >>> Pa = Particle('Pa', P, 1) + >>> Ac = Point('Ac') + >>> Ac.set_vel(N, 25 * N.y) + >>> I = outer(N.x, N.x) + >>> A = RigidBody('A', Ac, N, 20, (I, Ac)) + >>> linear_momentum(N, A, Pa) + 10*N.x + 500*N.y + + """ + + if not isinstance(frame, ReferenceFrame): + raise TypeError('Please specify a valid ReferenceFrame') + else: + linear_momentum_sys = Vector(0) + for e in body: + if isinstance(e, (RigidBody, Particle)): + linear_momentum_sys += e.linear_momentum(frame) + else: + raise TypeError('*body must have only Particle or RigidBody') + return linear_momentum_sys + + +def angular_momentum(point, frame, *body): + """Angular momentum of a system. + + Explanation + =========== + + This function returns the angular momentum of a system of Particle's and/or + RigidBody's. The angular momentum of such a system is equal to the vector + sum of the angular momentum of its constituents. Consider a system, S, + comprised of a rigid body, A, and a particle, P. The angular momentum of + the system, H, is equal to the vector sum of the angular momentum of the + particle, H1, and the angular momentum of the rigid body, H2, i.e. + + H = H1 + H2 + + Parameters + ========== + + point : Point + The point about which angular momentum of the system is desired. + frame : ReferenceFrame + The frame in which angular momentum is desired. + body1, body2, body3... : Particle and/or RigidBody + The body (or bodies) whose angular momentum is required. + + Examples + ======== + + >>> from sympy.physics.mechanics import Point, Particle, ReferenceFrame + >>> from sympy.physics.mechanics import RigidBody, outer, angular_momentum + >>> N = ReferenceFrame('N') + >>> O = Point('O') + >>> O.set_vel(N, 0 * N.x) + >>> P = O.locatenew('P', 1 * N.x) + >>> P.set_vel(N, 10 * N.x) + >>> Pa = Particle('Pa', P, 1) + >>> Ac = O.locatenew('Ac', 2 * N.y) + >>> Ac.set_vel(N, 5 * N.y) + >>> a = ReferenceFrame('a') + >>> a.set_ang_vel(N, 10 * N.z) + >>> I = outer(N.z, N.z) + >>> A = RigidBody('A', Ac, a, 20, (I, Ac)) + >>> angular_momentum(O, N, Pa, A) + 10*N.z + + """ + + if not isinstance(frame, ReferenceFrame): + raise TypeError('Please enter a valid ReferenceFrame') + if not isinstance(point, Point): + raise TypeError('Please specify a valid Point') + else: + angular_momentum_sys = Vector(0) + for e in body: + if isinstance(e, (RigidBody, Particle)): + angular_momentum_sys += e.angular_momentum(point, frame) + else: + raise TypeError('*body must have only Particle or RigidBody') + return angular_momentum_sys + + +def kinetic_energy(frame, *body): + """Kinetic energy of a multibody system. + + Explanation + =========== + + This function returns the kinetic energy of a system of Particle's and/or + RigidBody's. The kinetic energy of such a system is equal to the sum of + the kinetic energies of its constituents. Consider a system, S, comprising + a rigid body, A, and a particle, P. The kinetic energy of the system, T, + is equal to the vector sum of the kinetic energy of the particle, T1, and + the kinetic energy of the rigid body, T2, i.e. + + T = T1 + T2 + + Kinetic energy is a scalar. + + Parameters + ========== + + frame : ReferenceFrame + The frame in which the velocity or angular velocity of the body is + defined. + body1, body2, body3... : Particle and/or RigidBody + The body (or bodies) whose kinetic energy is required. + + Examples + ======== + + >>> from sympy.physics.mechanics import Point, Particle, ReferenceFrame + >>> from sympy.physics.mechanics import RigidBody, outer, kinetic_energy + >>> N = ReferenceFrame('N') + >>> O = Point('O') + >>> O.set_vel(N, 0 * N.x) + >>> P = O.locatenew('P', 1 * N.x) + >>> P.set_vel(N, 10 * N.x) + >>> Pa = Particle('Pa', P, 1) + >>> Ac = O.locatenew('Ac', 2 * N.y) + >>> Ac.set_vel(N, 5 * N.y) + >>> a = ReferenceFrame('a') + >>> a.set_ang_vel(N, 10 * N.z) + >>> I = outer(N.z, N.z) + >>> A = RigidBody('A', Ac, a, 20, (I, Ac)) + >>> kinetic_energy(N, Pa, A) + 350 + + """ + + if not isinstance(frame, ReferenceFrame): + raise TypeError('Please enter a valid ReferenceFrame') + ke_sys = S.Zero + for e in body: + if isinstance(e, (RigidBody, Particle)): + ke_sys += e.kinetic_energy(frame) + else: + raise TypeError('*body must have only Particle or RigidBody') + return ke_sys + + +def potential_energy(*body): + """Potential energy of a multibody system. + + Explanation + =========== + + This function returns the potential energy of a system of Particle's and/or + RigidBody's. The potential energy of such a system is equal to the sum of + the potential energy of its constituents. Consider a system, S, comprising + a rigid body, A, and a particle, P. The potential energy of the system, V, + is equal to the vector sum of the potential energy of the particle, V1, and + the potential energy of the rigid body, V2, i.e. + + V = V1 + V2 + + Potential energy is a scalar. + + Parameters + ========== + + body1, body2, body3... : Particle and/or RigidBody + The body (or bodies) whose potential energy is required. + + Examples + ======== + + >>> from sympy.physics.mechanics import Point, Particle, ReferenceFrame + >>> from sympy.physics.mechanics import RigidBody, outer, potential_energy + >>> from sympy import symbols + >>> M, m, g, h = symbols('M m g h') + >>> N = ReferenceFrame('N') + >>> O = Point('O') + >>> O.set_vel(N, 0 * N.x) + >>> P = O.locatenew('P', 1 * N.x) + >>> Pa = Particle('Pa', P, m) + >>> Ac = O.locatenew('Ac', 2 * N.y) + >>> a = ReferenceFrame('a') + >>> I = outer(N.z, N.z) + >>> A = RigidBody('A', Ac, a, M, (I, Ac)) + >>> Pa.potential_energy = m * g * h + >>> A.potential_energy = M * g * h + >>> potential_energy(Pa, A) + M*g*h + g*h*m + + """ + + pe_sys = S.Zero + for e in body: + if isinstance(e, (RigidBody, Particle)): + pe_sys += e.potential_energy + else: + raise TypeError('*body must have only Particle or RigidBody') + return pe_sys + + +def gravity(acceleration, *bodies): + from sympy.physics.mechanics.loads import gravity as _gravity + sympy_deprecation_warning( + """ + The gravity function has been moved. + Import it from "sympy.physics.mechanics.loads". + """, + deprecated_since_version="1.13", + active_deprecations_target="moved-mechanics-functions" + ) + return _gravity(acceleration, *bodies) + + +def center_of_mass(point, *bodies): + """ + Returns the position vector from the given point to the center of mass + of the given bodies(particles or rigidbodies). + + Example + ======= + + >>> from sympy import symbols, S + >>> from sympy.physics.vector import Point + >>> from sympy.physics.mechanics import Particle, ReferenceFrame, RigidBody, outer + >>> from sympy.physics.mechanics.functions import center_of_mass + >>> a = ReferenceFrame('a') + >>> m = symbols('m', real=True) + >>> p1 = Particle('p1', Point('p1_pt'), S(1)) + >>> p2 = Particle('p2', Point('p2_pt'), S(2)) + >>> p3 = Particle('p3', Point('p3_pt'), S(3)) + >>> p4 = Particle('p4', Point('p4_pt'), m) + >>> b_f = ReferenceFrame('b_f') + >>> b_cm = Point('b_cm') + >>> mb = symbols('mb') + >>> b = RigidBody('b', b_cm, b_f, mb, (outer(b_f.x, b_f.x), b_cm)) + >>> p2.point.set_pos(p1.point, a.x) + >>> p3.point.set_pos(p1.point, a.x + a.y) + >>> p4.point.set_pos(p1.point, a.y) + >>> b.masscenter.set_pos(p1.point, a.y + a.z) + >>> point_o=Point('o') + >>> point_o.set_pos(p1.point, center_of_mass(p1.point, p1, p2, p3, p4, b)) + >>> expr = 5/(m + mb + 6)*a.x + (m + mb + 3)/(m + mb + 6)*a.y + mb/(m + mb + 6)*a.z + >>> point_o.pos_from(p1.point) + 5/(m + mb + 6)*a.x + (m + mb + 3)/(m + mb + 6)*a.y + mb/(m + mb + 6)*a.z + + """ + if not bodies: + raise TypeError("No bodies(instances of Particle or Rigidbody) were passed.") + + total_mass = 0 + vec = Vector(0) + for i in bodies: + total_mass += i.mass + + masscenter = getattr(i, 'masscenter', None) + if masscenter is None: + masscenter = i.point + vec += i.mass*masscenter.pos_from(point) + + return vec/total_mass + + +def Lagrangian(frame, *body): + """Lagrangian of a multibody system. + + Explanation + =========== + + This function returns the Lagrangian of a system of Particle's and/or + RigidBody's. The Lagrangian of such a system is equal to the difference + between the kinetic energies and potential energies of its constituents. If + T and V are the kinetic and potential energies of a system then it's + Lagrangian, L, is defined as + + L = T - V + + The Lagrangian is a scalar. + + Parameters + ========== + + frame : ReferenceFrame + The frame in which the velocity or angular velocity of the body is + defined to determine the kinetic energy. + + body1, body2, body3... : Particle and/or RigidBody + The body (or bodies) whose Lagrangian is required. + + Examples + ======== + + >>> from sympy.physics.mechanics import Point, Particle, ReferenceFrame + >>> from sympy.physics.mechanics import RigidBody, outer, Lagrangian + >>> from sympy import symbols + >>> M, m, g, h = symbols('M m g h') + >>> N = ReferenceFrame('N') + >>> O = Point('O') + >>> O.set_vel(N, 0 * N.x) + >>> P = O.locatenew('P', 1 * N.x) + >>> P.set_vel(N, 10 * N.x) + >>> Pa = Particle('Pa', P, 1) + >>> Ac = O.locatenew('Ac', 2 * N.y) + >>> Ac.set_vel(N, 5 * N.y) + >>> a = ReferenceFrame('a') + >>> a.set_ang_vel(N, 10 * N.z) + >>> I = outer(N.z, N.z) + >>> A = RigidBody('A', Ac, a, 20, (I, Ac)) + >>> Pa.potential_energy = m * g * h + >>> A.potential_energy = M * g * h + >>> Lagrangian(N, Pa, A) + -M*g*h - g*h*m + 350 + + """ + + if not isinstance(frame, ReferenceFrame): + raise TypeError('Please supply a valid ReferenceFrame') + for e in body: + if not isinstance(e, (RigidBody, Particle)): + raise TypeError('*body must have only Particle or RigidBody') + return kinetic_energy(frame, *body) - potential_energy(*body) + + +def find_dynamicsymbols(expression, exclude=None, reference_frame=None): + """Find all dynamicsymbols in expression. + + Explanation + =========== + + If the optional ``exclude`` kwarg is used, only dynamicsymbols + not in the iterable ``exclude`` are returned. + If we intend to apply this function on a vector, the optional + ``reference_frame`` is also used to inform about the corresponding frame + with respect to which the dynamic symbols of the given vector is to be + determined. + + Parameters + ========== + + expression : SymPy expression + + exclude : iterable of dynamicsymbols, optional + + reference_frame : ReferenceFrame, optional + The frame with respect to which the dynamic symbols of the + given vector is to be determined. + + Examples + ======== + + >>> from sympy.physics.mechanics import dynamicsymbols, find_dynamicsymbols + >>> from sympy.physics.mechanics import ReferenceFrame + >>> x, y = dynamicsymbols('x, y') + >>> expr = x + x.diff()*y + >>> find_dynamicsymbols(expr) + {x(t), y(t), Derivative(x(t), t)} + >>> find_dynamicsymbols(expr, exclude=[x, y]) + {Derivative(x(t), t)} + >>> a, b, c = dynamicsymbols('a, b, c') + >>> A = ReferenceFrame('A') + >>> v = a * A.x + b * A.y + c * A.z + >>> find_dynamicsymbols(v, reference_frame=A) + {a(t), b(t), c(t)} + + """ + t_set = {dynamicsymbols._t} + if exclude: + if iterable(exclude): + exclude_set = set(exclude) + else: + raise TypeError("exclude kwarg must be iterable") + else: + exclude_set = set() + if isinstance(expression, Vector): + if reference_frame is None: + raise ValueError("You must provide reference_frame when passing a " + "vector expression, got %s." % reference_frame) + else: + expression = expression.to_matrix(reference_frame) + return {i for i in expression.atoms(AppliedUndef, Derivative) if + i.free_symbols == t_set} - exclude_set + + +def msubs(expr, *sub_dicts, smart=False, **kwargs): + """A custom subs for use on expressions derived in physics.mechanics. + + Traverses the expression tree once, performing the subs found in sub_dicts. + Terms inside ``Derivative`` expressions are ignored: + + Examples + ======== + + >>> from sympy.physics.mechanics import dynamicsymbols, msubs + >>> x = dynamicsymbols('x') + >>> msubs(x.diff() + x, {x: 1}) + Derivative(x(t), t) + 1 + + Note that sub_dicts can be a single dictionary, or several dictionaries: + + >>> x, y, z = dynamicsymbols('x, y, z') + >>> sub1 = {x: 1, y: 2} + >>> sub2 = {z: 3, x.diff(): 4} + >>> msubs(x.diff() + x + y + z, sub1, sub2) + 10 + + If smart=True (default False), also checks for conditions that may result + in ``nan``, but if simplified would yield a valid expression. For example: + + >>> from sympy import sin, tan + >>> (sin(x)/tan(x)).subs(x, 0) + nan + >>> msubs(sin(x)/tan(x), {x: 0}, smart=True) + 1 + + It does this by first replacing all ``tan`` with ``sin/cos``. Then each + node is traversed. If the node is a fraction, subs is first evaluated on + the denominator. If this results in 0, simplification of the entire + fraction is attempted. Using this selective simplification, only + subexpressions that result in 1/0 are targeted, resulting in faster + performance. + + """ + + sub_dict = dict_merge(*sub_dicts) + if smart: + func = _smart_subs + elif hasattr(expr, 'msubs'): + return expr.msubs(sub_dict) + else: + func = lambda expr, sub_dict: _crawl(expr, _sub_func, sub_dict) + if isinstance(expr, (Matrix, Vector, Dyadic)): + return expr.applyfunc(lambda x: func(x, sub_dict)) + else: + return func(expr, sub_dict) + + +def _crawl(expr, func, *args, **kwargs): + """Crawl the expression tree, and apply func to every node.""" + val = func(expr, *args, **kwargs) + if val is not None: + return val + new_args = (_crawl(arg, func, *args, **kwargs) for arg in expr.args) + return expr.func(*new_args) + + +def _sub_func(expr, sub_dict): + """Perform direct matching substitution, ignoring derivatives.""" + if expr in sub_dict: + return sub_dict[expr] + elif not expr.args or expr.is_Derivative: + return expr + + +def _tan_repl_func(expr): + """Replace tan with sin/cos.""" + if isinstance(expr, tan): + return sin(*expr.args) / cos(*expr.args) + elif not expr.args or expr.is_Derivative: + return expr + + +def _smart_subs(expr, sub_dict): + """Performs subs, checking for conditions that may result in `nan` or + `oo`, and attempts to simplify them out. + + The expression tree is traversed twice, and the following steps are + performed on each expression node: + - First traverse: + Replace all `tan` with `sin/cos`. + - Second traverse: + If node is a fraction, check if the denominator evaluates to 0. + If so, attempt to simplify it out. Then if node is in sub_dict, + sub in the corresponding value. + + """ + expr = _crawl(expr, _tan_repl_func) + + def _recurser(expr, sub_dict): + # Decompose the expression into num, den + num, den = _fraction_decomp(expr) + if den != 1: + # If there is a non trivial denominator, we need to handle it + denom_subbed = _recurser(den, sub_dict) + if denom_subbed.evalf() == 0: + # If denom is 0 after this, attempt to simplify the bad expr + expr = simplify(expr) + else: + # Expression won't result in nan, find numerator + num_subbed = _recurser(num, sub_dict) + return num_subbed / denom_subbed + # We have to crawl the tree manually, because `expr` may have been + # modified in the simplify step. First, perform subs as normal: + val = _sub_func(expr, sub_dict) + if val is not None: + return val + new_args = (_recurser(arg, sub_dict) for arg in expr.args) + return expr.func(*new_args) + return _recurser(expr, sub_dict) + + +def _fraction_decomp(expr): + """Return num, den such that expr = num/den.""" + if not isinstance(expr, Mul): + return expr, 1 + num = [] + den = [] + for a in expr.args: + if a.is_Pow and a.args[1] < 0: + den.append(1 / a) + else: + num.append(a) + if not den: + return expr, 1 + num = Mul(*num) + den = Mul(*den) + return num, den + + +def _f_list_parser(fl, ref_frame): + """Parses the provided forcelist composed of items + of the form (obj, force). + Returns a tuple containing: + vel_list: The velocity (ang_vel for Frames, vel for Points) in + the provided reference frame. + f_list: The forces. + + Used internally in the KanesMethod and LagrangesMethod classes. + + """ + def flist_iter(): + for pair in fl: + obj, force = pair + if isinstance(obj, ReferenceFrame): + yield obj.ang_vel_in(ref_frame), force + elif isinstance(obj, Point): + yield obj.vel(ref_frame), force + else: + raise TypeError('First entry in each forcelist pair must ' + 'be a point or frame.') + + if not fl: + vel_list, f_list = (), () + else: + unzip = lambda l: list(zip(*l)) if l[0] else [(), ()] + vel_list, f_list = unzip(list(flist_iter())) + return vel_list, f_list + + +def _validate_coordinates(coordinates=None, speeds=None, check_duplicates=True, + is_dynamicsymbols=True, u_auxiliary=None): + """Validate the generalized coordinates and generalized speeds. + + Parameters + ========== + coordinates : iterable, optional + Generalized coordinates to be validated. + speeds : iterable, optional + Generalized speeds to be validated. + check_duplicates : bool, optional + Checks if there are duplicates in the generalized coordinates and + generalized speeds. If so it will raise a ValueError. The default is + True. + is_dynamicsymbols : iterable, optional + Checks if all the generalized coordinates and generalized speeds are + dynamicsymbols. If any is not a dynamicsymbol, a ValueError will be + raised. The default is True. + u_auxiliary : iterable, optional + Auxiliary generalized speeds to be validated. + + """ + t_set = {dynamicsymbols._t} + # Convert input to iterables + if coordinates is None: + coordinates = [] + elif not iterable(coordinates): + coordinates = [coordinates] + if speeds is None: + speeds = [] + elif not iterable(speeds): + speeds = [speeds] + if u_auxiliary is None: + u_auxiliary = [] + elif not iterable(u_auxiliary): + u_auxiliary = [u_auxiliary] + + msgs = [] + if check_duplicates: # Check for duplicates + seen = set() + coord_duplicates = {x for x in coordinates if x in seen or seen.add(x)} + seen = set() + speed_duplicates = {x for x in speeds if x in seen or seen.add(x)} + seen = set() + aux_duplicates = {x for x in u_auxiliary if x in seen or seen.add(x)} + overlap_coords = set(coordinates).intersection(speeds) + overlap_aux = set(coordinates).union(speeds).intersection(u_auxiliary) + if coord_duplicates: + msgs.append(f'The generalized coordinates {coord_duplicates} are ' + f'duplicated, all generalized coordinates should be ' + f'unique.') + if speed_duplicates: + msgs.append(f'The generalized speeds {speed_duplicates} are ' + f'duplicated, all generalized speeds should be unique.') + if aux_duplicates: + msgs.append(f'The auxiliary speeds {aux_duplicates} are duplicated,' + f' all auxiliary speeds should be unique.') + if overlap_coords: + msgs.append(f'{overlap_coords} are defined as both generalized ' + f'coordinates and generalized speeds.') + if overlap_aux: + msgs.append(f'The auxiliary speeds {overlap_aux} are also defined ' + f'as generalized coordinates or generalized speeds.') + if is_dynamicsymbols: # Check whether all coordinates are dynamicsymbols + for coordinate in coordinates: + if not (isinstance(coordinate, (AppliedUndef, Derivative)) and + coordinate.free_symbols == t_set): + msgs.append(f'Generalized coordinate "{coordinate}" is not a ' + f'dynamicsymbol.') + for speed in speeds: + if not (isinstance(speed, (AppliedUndef, Derivative)) and + speed.free_symbols == t_set): + msgs.append( + f'Generalized speed "{speed}" is not a dynamicsymbol.') + for aux in u_auxiliary: + if not (isinstance(aux, (AppliedUndef, Derivative)) and + aux.free_symbols == t_set): + msgs.append( + f'Auxiliary speed "{aux}" is not a dynamicsymbol.') + if msgs: + raise ValueError('\n'.join(msgs)) + + +def _parse_linear_solver(linear_solver): + """Helper function to retrieve a specified linear solver.""" + if callable(linear_solver): + return linear_solver + return lambda A, b: Matrix.solve(A, b, method=linear_solver) diff --git a/lib/python3.12/site-packages/sympy/physics/mechanics/inertia.py b/lib/python3.12/site-packages/sympy/physics/mechanics/inertia.py new file mode 100644 index 0000000000000000000000000000000000000000..d2fe37c7f9f39c692f98c8bf038a73326e171dd7 --- /dev/null +++ b/lib/python3.12/site-packages/sympy/physics/mechanics/inertia.py @@ -0,0 +1,197 @@ +from sympy import sympify +from sympy.physics.vector import Point, Dyadic, ReferenceFrame, outer +from collections import namedtuple + +__all__ = ['inertia', 'inertia_of_point_mass', 'Inertia'] + + +def inertia(frame, ixx, iyy, izz, ixy=0, iyz=0, izx=0): + """Simple way to create inertia Dyadic object. + + Explanation + =========== + + Creates an inertia Dyadic based on the given tensor values and a body-fixed + reference frame. + + Parameters + ========== + + frame : ReferenceFrame + The frame the inertia is defined in. + ixx : Sympifyable + The xx element in the inertia dyadic. + iyy : Sympifyable + The yy element in the inertia dyadic. + izz : Sympifyable + The zz element in the inertia dyadic. + ixy : Sympifyable + The xy element in the inertia dyadic. + iyz : Sympifyable + The yz element in the inertia dyadic. + izx : Sympifyable + The zx element in the inertia dyadic. + + Examples + ======== + + >>> from sympy.physics.mechanics import ReferenceFrame, inertia + >>> N = ReferenceFrame('N') + >>> inertia(N, 1, 2, 3) + (N.x|N.x) + 2*(N.y|N.y) + 3*(N.z|N.z) + + """ + + if not isinstance(frame, ReferenceFrame): + raise TypeError('Need to define the inertia in a frame') + ixx, iyy, izz = sympify(ixx), sympify(iyy), sympify(izz) + ixy, iyz, izx = sympify(ixy), sympify(iyz), sympify(izx) + return (ixx*outer(frame.x, frame.x) + ixy*outer(frame.x, frame.y) + + izx*outer(frame.x, frame.z) + ixy*outer(frame.y, frame.x) + + iyy*outer(frame.y, frame.y) + iyz*outer(frame.y, frame.z) + + izx*outer(frame.z, frame.x) + iyz*outer(frame.z, frame.y) + + izz*outer(frame.z, frame.z)) + + +def inertia_of_point_mass(mass, pos_vec, frame): + """Inertia dyadic of a point mass relative to point O. + + Parameters + ========== + + mass : Sympifyable + Mass of the point mass + pos_vec : Vector + Position from point O to point mass + frame : ReferenceFrame + Reference frame to express the dyadic in + + Examples + ======== + + >>> from sympy import symbols + >>> from sympy.physics.mechanics import ReferenceFrame, inertia_of_point_mass + >>> N = ReferenceFrame('N') + >>> r, m = symbols('r m') + >>> px = r * N.x + >>> inertia_of_point_mass(m, px, N) + m*r**2*(N.y|N.y) + m*r**2*(N.z|N.z) + + """ + + return mass*( + (outer(frame.x, frame.x) + + outer(frame.y, frame.y) + + outer(frame.z, frame.z)) * + (pos_vec.dot(pos_vec)) - outer(pos_vec, pos_vec)) + + +class Inertia(namedtuple('Inertia', ['dyadic', 'point'])): + """Inertia object consisting of a Dyadic and a Point of reference. + + Explanation + =========== + + This is a simple class to store the Point and Dyadic, belonging to an + inertia. + + Attributes + ========== + + dyadic : Dyadic + The dyadic of the inertia. + point : Point + The reference point of the inertia. + + Examples + ======== + + >>> from sympy.physics.mechanics import ReferenceFrame, Point, Inertia + >>> N = ReferenceFrame('N') + >>> Po = Point('Po') + >>> Inertia(N.x.outer(N.x) + N.y.outer(N.y) + N.z.outer(N.z), Po) + ((N.x|N.x) + (N.y|N.y) + (N.z|N.z), Po) + + In the example above the Dyadic was created manually, one can however also + use the ``inertia`` function for this or the class method ``from_tensor`` as + shown below. + + >>> Inertia.from_inertia_scalars(Po, N, 1, 1, 1) + ((N.x|N.x) + (N.y|N.y) + (N.z|N.z), Po) + + """ + def __new__(cls, dyadic, point): + # Switch order if given in the wrong order + if isinstance(dyadic, Point) and isinstance(point, Dyadic): + point, dyadic = dyadic, point + if not isinstance(point, Point): + raise TypeError('Reference point should be of type Point') + if not isinstance(dyadic, Dyadic): + raise TypeError('Inertia value should be expressed as a Dyadic') + return super().__new__(cls, dyadic, point) + + @classmethod + def from_inertia_scalars(cls, point, frame, ixx, iyy, izz, ixy=0, iyz=0, + izx=0): + """Simple way to create an Inertia object based on the tensor values. + + Explanation + =========== + + This class method uses the :func`~.inertia` to create the Dyadic based + on the tensor values. + + Parameters + ========== + + point : Point + The reference point of the inertia. + frame : ReferenceFrame + The frame the inertia is defined in. + ixx : Sympifyable + The xx element in the inertia dyadic. + iyy : Sympifyable + The yy element in the inertia dyadic. + izz : Sympifyable + The zz element in the inertia dyadic. + ixy : Sympifyable + The xy element in the inertia dyadic. + iyz : Sympifyable + The yz element in the inertia dyadic. + izx : Sympifyable + The zx element in the inertia dyadic. + + Examples + ======== + + >>> from sympy import symbols + >>> from sympy.physics.mechanics import ReferenceFrame, Point, Inertia + >>> ixx, iyy, izz, ixy, iyz, izx = symbols('ixx iyy izz ixy iyz izx') + >>> N = ReferenceFrame('N') + >>> P = Point('P') + >>> I = Inertia.from_inertia_scalars(P, N, ixx, iyy, izz, ixy, iyz, izx) + + The tensor values can easily be seen when converting the dyadic to a + matrix. + + >>> I.dyadic.to_matrix(N) + Matrix([ + [ixx, ixy, izx], + [ixy, iyy, iyz], + [izx, iyz, izz]]) + + """ + return cls(inertia(frame, ixx, iyy, izz, ixy, iyz, izx), point) + + def __add__(self, other): + raise TypeError(f"unsupported operand type(s) for +: " + f"'{self.__class__.__name__}' and " + f"'{other.__class__.__name__}'") + + def __mul__(self, other): + raise TypeError(f"unsupported operand type(s) for *: " + f"'{self.__class__.__name__}' and " + f"'{other.__class__.__name__}'") + + __radd__ = __add__ + __rmul__ = __mul__ diff --git a/lib/python3.12/site-packages/sympy/physics/mechanics/lagrange.py b/lib/python3.12/site-packages/sympy/physics/mechanics/lagrange.py new file mode 100644 index 0000000000000000000000000000000000000000..282176a404f77762abc3ee8c6a575519b2de1f02 --- /dev/null +++ b/lib/python3.12/site-packages/sympy/physics/mechanics/lagrange.py @@ -0,0 +1,512 @@ +from sympy import diff, zeros, Matrix, eye, sympify +from sympy.core.sorting import default_sort_key +from sympy.physics.vector import dynamicsymbols, ReferenceFrame +from sympy.physics.mechanics.method import _Methods +from sympy.physics.mechanics.functions import ( + find_dynamicsymbols, msubs, _f_list_parser, _validate_coordinates) +from sympy.physics.mechanics.linearize import Linearizer +from sympy.utilities.iterables import iterable + +__all__ = ['LagrangesMethod'] + + +class LagrangesMethod(_Methods): + """Lagrange's method object. + + Explanation + =========== + + This object generates the equations of motion in a two step procedure. The + first step involves the initialization of LagrangesMethod by supplying the + Lagrangian and the generalized coordinates, at the bare minimum. If there + are any constraint equations, they can be supplied as keyword arguments. + The Lagrange multipliers are automatically generated and are equal in + number to the constraint equations. Similarly any non-conservative forces + can be supplied in an iterable (as described below and also shown in the + example) along with a ReferenceFrame. This is also discussed further in the + __init__ method. + + Attributes + ========== + + q, u : Matrix + Matrices of the generalized coordinates and speeds + loads : iterable + Iterable of (Point, vector) or (ReferenceFrame, vector) tuples + describing the forces on the system. + bodies : iterable + Iterable containing the rigid bodies and particles of the system. + mass_matrix : Matrix + The system's mass matrix + forcing : Matrix + The system's forcing vector + mass_matrix_full : Matrix + The "mass matrix" for the qdot's, qdoubledot's, and the + lagrange multipliers (lam) + forcing_full : Matrix + The forcing vector for the qdot's, qdoubledot's and + lagrange multipliers (lam) + + Examples + ======== + + This is a simple example for a one degree of freedom translational + spring-mass-damper. + + In this example, we first need to do the kinematics. + This involves creating generalized coordinates and their derivatives. + Then we create a point and set its velocity in a frame. + + >>> from sympy.physics.mechanics import LagrangesMethod, Lagrangian + >>> from sympy.physics.mechanics import ReferenceFrame, Particle, Point + >>> from sympy.physics.mechanics import dynamicsymbols + >>> from sympy import symbols + >>> q = dynamicsymbols('q') + >>> qd = dynamicsymbols('q', 1) + >>> m, k, b = symbols('m k b') + >>> N = ReferenceFrame('N') + >>> P = Point('P') + >>> P.set_vel(N, qd * N.x) + + We need to then prepare the information as required by LagrangesMethod to + generate equations of motion. + First we create the Particle, which has a point attached to it. + Following this the lagrangian is created from the kinetic and potential + energies. + Then, an iterable of nonconservative forces/torques must be constructed, + where each item is a (Point, Vector) or (ReferenceFrame, Vector) tuple, + with the Vectors representing the nonconservative forces or torques. + + >>> Pa = Particle('Pa', P, m) + >>> Pa.potential_energy = k * q**2 / 2.0 + >>> L = Lagrangian(N, Pa) + >>> fl = [(P, -b * qd * N.x)] + + Finally we can generate the equations of motion. + First we create the LagrangesMethod object. To do this one must supply + the Lagrangian, and the generalized coordinates. The constraint equations, + the forcelist, and the inertial frame may also be provided, if relevant. + Next we generate Lagrange's equations of motion, such that: + Lagrange's equations of motion = 0. + We have the equations of motion at this point. + + >>> l = LagrangesMethod(L, [q], forcelist = fl, frame = N) + >>> print(l.form_lagranges_equations()) + Matrix([[b*Derivative(q(t), t) + 1.0*k*q(t) + m*Derivative(q(t), (t, 2))]]) + + We can also solve for the states using the 'rhs' method. + + >>> print(l.rhs()) + Matrix([[Derivative(q(t), t)], [(-b*Derivative(q(t), t) - 1.0*k*q(t))/m]]) + + Please refer to the docstrings on each method for more details. + """ + + def __init__(self, Lagrangian, qs, forcelist=None, bodies=None, frame=None, + hol_coneqs=None, nonhol_coneqs=None): + """Supply the following for the initialization of LagrangesMethod. + + Lagrangian : Sympifyable + + qs : array_like + The generalized coordinates + + hol_coneqs : array_like, optional + The holonomic constraint equations + + nonhol_coneqs : array_like, optional + The nonholonomic constraint equations + + forcelist : iterable, optional + Takes an iterable of (Point, Vector) or (ReferenceFrame, Vector) + tuples which represent the force at a point or torque on a frame. + This feature is primarily to account for the nonconservative forces + and/or moments. + + bodies : iterable, optional + Takes an iterable containing the rigid bodies and particles of the + system. + + frame : ReferenceFrame, optional + Supply the inertial frame. This is used to determine the + generalized forces due to non-conservative forces. + """ + + self._L = Matrix([sympify(Lagrangian)]) + self.eom = None + self._m_cd = Matrix() # Mass Matrix of differentiated coneqs + self._m_d = Matrix() # Mass Matrix of dynamic equations + self._f_cd = Matrix() # Forcing part of the diff coneqs + self._f_d = Matrix() # Forcing part of the dynamic equations + self.lam_coeffs = Matrix() # The coeffecients of the multipliers + + forcelist = forcelist if forcelist else [] + if not iterable(forcelist): + raise TypeError('Force pairs must be supplied in an iterable.') + self._forcelist = forcelist + if frame and not isinstance(frame, ReferenceFrame): + raise TypeError('frame must be a valid ReferenceFrame') + self._bodies = bodies + self.inertial = frame + + self.lam_vec = Matrix() + + self._term1 = Matrix() + self._term2 = Matrix() + self._term3 = Matrix() + self._term4 = Matrix() + + # Creating the qs, qdots and qdoubledots + if not iterable(qs): + raise TypeError('Generalized coordinates must be an iterable') + self._q = Matrix(qs) + self._qdots = self.q.diff(dynamicsymbols._t) + self._qdoubledots = self._qdots.diff(dynamicsymbols._t) + _validate_coordinates(self.q) + + mat_build = lambda x: Matrix(x) if x else Matrix() + hol_coneqs = mat_build(hol_coneqs) + nonhol_coneqs = mat_build(nonhol_coneqs) + self.coneqs = Matrix([hol_coneqs.diff(dynamicsymbols._t), + nonhol_coneqs]) + self._hol_coneqs = hol_coneqs + + def form_lagranges_equations(self): + """Method to form Lagrange's equations of motion. + + Returns a vector of equations of motion using Lagrange's equations of + the second kind. + """ + + qds = self._qdots + qdd_zero = dict.fromkeys(self._qdoubledots, 0) + n = len(self.q) + + # Internally we represent the EOM as four terms: + # EOM = term1 - term2 - term3 - term4 = 0 + + # First term + self._term1 = self._L.jacobian(qds) + self._term1 = self._term1.diff(dynamicsymbols._t).T + + # Second term + self._term2 = self._L.jacobian(self.q).T + + # Third term + if self.coneqs: + coneqs = self.coneqs + m = len(coneqs) + # Creating the multipliers + self.lam_vec = Matrix(dynamicsymbols('lam1:' + str(m + 1))) + self.lam_coeffs = -coneqs.jacobian(qds) + self._term3 = self.lam_coeffs.T * self.lam_vec + # Extracting the coeffecients of the qdds from the diff coneqs + diffconeqs = coneqs.diff(dynamicsymbols._t) + self._m_cd = diffconeqs.jacobian(self._qdoubledots) + # The remaining terms i.e. the 'forcing' terms in diff coneqs + self._f_cd = -diffconeqs.subs(qdd_zero) + else: + self._term3 = zeros(n, 1) + + # Fourth term + if self.forcelist: + N = self.inertial + self._term4 = zeros(n, 1) + for i, qd in enumerate(qds): + flist = zip(*_f_list_parser(self.forcelist, N)) + self._term4[i] = sum(v.diff(qd, N).dot(f) for (v, f) in flist) + else: + self._term4 = zeros(n, 1) + + # Form the dynamic mass and forcing matrices + without_lam = self._term1 - self._term2 - self._term4 + self._m_d = without_lam.jacobian(self._qdoubledots) + self._f_d = -without_lam.subs(qdd_zero) + + # Form the EOM + self.eom = without_lam - self._term3 + return self.eom + + def _form_eoms(self): + return self.form_lagranges_equations() + + @property + def mass_matrix(self): + """Returns the mass matrix, which is augmented by the Lagrange + multipliers, if necessary. + + Explanation + =========== + + If the system is described by 'n' generalized coordinates and there are + no constraint equations then an n X n matrix is returned. + + If there are 'n' generalized coordinates and 'm' constraint equations + have been supplied during initialization then an n X (n+m) matrix is + returned. The (n + m - 1)th and (n + m)th columns contain the + coefficients of the Lagrange multipliers. + """ + + if self.eom is None: + raise ValueError('Need to compute the equations of motion first') + if self.coneqs: + return (self._m_d).row_join(self.lam_coeffs.T) + else: + return self._m_d + + @property + def mass_matrix_full(self): + """Augments the coefficients of qdots to the mass_matrix.""" + + if self.eom is None: + raise ValueError('Need to compute the equations of motion first') + n = len(self.q) + m = len(self.coneqs) + row1 = eye(n).row_join(zeros(n, n + m)) + row2 = zeros(n, n).row_join(self.mass_matrix) + if self.coneqs: + row3 = zeros(m, n).row_join(self._m_cd).row_join(zeros(m, m)) + return row1.col_join(row2).col_join(row3) + else: + return row1.col_join(row2) + + @property + def forcing(self): + """Returns the forcing vector from 'lagranges_equations' method.""" + + if self.eom is None: + raise ValueError('Need to compute the equations of motion first') + return self._f_d + + @property + def forcing_full(self): + """Augments qdots to the forcing vector above.""" + + if self.eom is None: + raise ValueError('Need to compute the equations of motion first') + if self.coneqs: + return self._qdots.col_join(self.forcing).col_join(self._f_cd) + else: + return self._qdots.col_join(self.forcing) + + def to_linearizer(self, q_ind=None, qd_ind=None, q_dep=None, qd_dep=None, + linear_solver='LU'): + """Returns an instance of the Linearizer class, initiated from the data + in the LagrangesMethod class. This may be more desirable than using the + linearize class method, as the Linearizer object will allow more + efficient recalculation (i.e. about varying operating points). + + Parameters + ========== + + q_ind, qd_ind : array_like, optional + The independent generalized coordinates and speeds. + q_dep, qd_dep : array_like, optional + The dependent generalized coordinates and speeds. + linear_solver : str, callable + Method used to solve the several symbolic linear systems of the + form ``A*x=b`` in the linearization process. If a string is + supplied, it should be a valid method that can be used with the + :meth:`sympy.matrices.matrixbase.MatrixBase.solve`. If a callable is + supplied, it should have the format ``x = f(A, b)``, where it + solves the equations and returns the solution. The default is + ``'LU'`` which corresponds to SymPy's ``A.LUsolve(b)``. + ``LUsolve()`` is fast to compute but will often result in + divide-by-zero and thus ``nan`` results. + + Returns + ======= + Linearizer + An instantiated + :class:`sympy.physics.mechanics.linearize.Linearizer`. + + """ + + # Compose vectors + t = dynamicsymbols._t + q = self.q + u = self._qdots + ud = u.diff(t) + # Get vector of lagrange multipliers + lams = self.lam_vec + + mat_build = lambda x: Matrix(x) if x else Matrix() + q_i = mat_build(q_ind) + q_d = mat_build(q_dep) + u_i = mat_build(qd_ind) + u_d = mat_build(qd_dep) + + # Compose general form equations + f_c = self._hol_coneqs + f_v = self.coneqs + f_a = f_v.diff(t) + f_0 = u + f_1 = -u + f_2 = self._term1 + f_3 = -(self._term2 + self._term4) + f_4 = -self._term3 + + # Check that there are an appropriate number of independent and + # dependent coordinates + if len(q_d) != len(f_c) or len(u_d) != len(f_v): + raise ValueError(("Must supply {:} dependent coordinates, and " + + "{:} dependent speeds").format(len(f_c), len(f_v))) + if set(Matrix([q_i, q_d])) != set(q): + raise ValueError("Must partition q into q_ind and q_dep, with " + + "no extra or missing symbols.") + if set(Matrix([u_i, u_d])) != set(u): + raise ValueError("Must partition qd into qd_ind and qd_dep, " + + "with no extra or missing symbols.") + + # Find all other dynamic symbols, forming the forcing vector r. + # Sort r to make it canonical. + insyms = set(Matrix([q, u, ud, lams])) + r = list(find_dynamicsymbols(f_3, insyms)) + r.sort(key=default_sort_key) + # Check for any derivatives of variables in r that are also found in r. + for i in r: + if diff(i, dynamicsymbols._t) in r: + raise ValueError('Cannot have derivatives of specified \ + quantities when linearizing forcing terms.') + + return Linearizer(f_0, f_1, f_2, f_3, f_4, f_c, f_v, f_a, q, u, q_i, + q_d, u_i, u_d, r, lams, linear_solver=linear_solver) + + def linearize(self, q_ind=None, qd_ind=None, q_dep=None, qd_dep=None, + linear_solver='LU', **kwargs): + """Linearize the equations of motion about a symbolic operating point. + + Parameters + ========== + linear_solver : str, callable + Method used to solve the several symbolic linear systems of the + form ``A*x=b`` in the linearization process. If a string is + supplied, it should be a valid method that can be used with the + :meth:`sympy.matrices.matrixbase.MatrixBase.solve`. If a callable is + supplied, it should have the format ``x = f(A, b)``, where it + solves the equations and returns the solution. The default is + ``'LU'`` which corresponds to SymPy's ``A.LUsolve(b)``. + ``LUsolve()`` is fast to compute but will often result in + divide-by-zero and thus ``nan`` results. + **kwargs + Extra keyword arguments are passed to + :meth:`sympy.physics.mechanics.linearize.Linearizer.linearize`. + + Explanation + =========== + + If kwarg A_and_B is False (default), returns M, A, B, r for the + linearized form, M*[q', u']^T = A*[q_ind, u_ind]^T + B*r. + + If kwarg A_and_B is True, returns A, B, r for the linearized form + dx = A*x + B*r, where x = [q_ind, u_ind]^T. Note that this is + computationally intensive if there are many symbolic parameters. For + this reason, it may be more desirable to use the default A_and_B=False, + returning M, A, and B. Values may then be substituted in to these + matrices, and the state space form found as + A = P.T*M.inv()*A, B = P.T*M.inv()*B, where P = Linearizer.perm_mat. + + In both cases, r is found as all dynamicsymbols in the equations of + motion that are not part of q, u, q', or u'. They are sorted in + canonical form. + + The operating points may be also entered using the ``op_point`` kwarg. + This takes a dictionary of {symbol: value}, or a an iterable of such + dictionaries. The values may be numeric or symbolic. The more values + you can specify beforehand, the faster this computation will run. + + For more documentation, please see the ``Linearizer`` class.""" + + linearizer = self.to_linearizer(q_ind, qd_ind, q_dep, qd_dep, + linear_solver=linear_solver) + result = linearizer.linearize(**kwargs) + return result + (linearizer.r,) + + def solve_multipliers(self, op_point=None, sol_type='dict'): + """Solves for the values of the lagrange multipliers symbolically at + the specified operating point. + + Parameters + ========== + + op_point : dict or iterable of dicts, optional + Point at which to solve at. The operating point is specified as + a dictionary or iterable of dictionaries of {symbol: value}. The + value may be numeric or symbolic itself. + + sol_type : str, optional + Solution return type. Valid options are: + - 'dict': A dict of {symbol : value} (default) + - 'Matrix': An ordered column matrix of the solution + """ + + # Determine number of multipliers + k = len(self.lam_vec) + if k == 0: + raise ValueError("System has no lagrange multipliers to solve for.") + # Compose dict of operating conditions + if isinstance(op_point, dict): + op_point_dict = op_point + elif iterable(op_point): + op_point_dict = {} + for op in op_point: + op_point_dict.update(op) + elif op_point is None: + op_point_dict = {} + else: + raise TypeError("op_point must be either a dictionary or an " + "iterable of dictionaries.") + # Compose the system to be solved + mass_matrix = self.mass_matrix.col_join(-self.lam_coeffs.row_join( + zeros(k, k))) + force_matrix = self.forcing.col_join(self._f_cd) + # Sub in the operating point + mass_matrix = msubs(mass_matrix, op_point_dict) + force_matrix = msubs(force_matrix, op_point_dict) + # Solve for the multipliers + sol_list = mass_matrix.LUsolve(-force_matrix)[-k:] + if sol_type == 'dict': + return dict(zip(self.lam_vec, sol_list)) + elif sol_type == 'Matrix': + return Matrix(sol_list) + else: + raise ValueError("Unknown sol_type {:}.".format(sol_type)) + + def rhs(self, inv_method=None, **kwargs): + """Returns equations that can be solved numerically. + + Parameters + ========== + + inv_method : str + The specific sympy inverse matrix calculation method to use. For a + list of valid methods, see + :meth:`~sympy.matrices.matrixbase.MatrixBase.inv` + """ + + if inv_method is None: + self._rhs = self.mass_matrix_full.LUsolve(self.forcing_full) + else: + self._rhs = (self.mass_matrix_full.inv(inv_method, + try_block_diag=True) * self.forcing_full) + return self._rhs + + @property + def q(self): + return self._q + + @property + def u(self): + return self._qdots + + @property + def bodies(self): + return self._bodies + + @property + def forcelist(self): + return self._forcelist + + @property + def loads(self): + return self._forcelist diff --git a/lib/python3.12/site-packages/sympy/physics/mechanics/loads.py b/lib/python3.12/site-packages/sympy/physics/mechanics/loads.py new file mode 100644 index 0000000000000000000000000000000000000000..3b9db763ffd6f99905e9d17fdc07f4171de4801b --- /dev/null +++ b/lib/python3.12/site-packages/sympy/physics/mechanics/loads.py @@ -0,0 +1,177 @@ +from abc import ABC +from collections import namedtuple +from sympy.physics.mechanics.body_base import BodyBase +from sympy.physics.vector import Vector, ReferenceFrame, Point + +__all__ = ['LoadBase', 'Force', 'Torque'] + + +class LoadBase(ABC, namedtuple('LoadBase', ['location', 'vector'])): + """Abstract base class for the various loading types.""" + + def __add__(self, other): + raise TypeError(f"unsupported operand type(s) for +: " + f"'{self.__class__.__name__}' and " + f"'{other.__class__.__name__}'") + + def __mul__(self, other): + raise TypeError(f"unsupported operand type(s) for *: " + f"'{self.__class__.__name__}' and " + f"'{other.__class__.__name__}'") + + __radd__ = __add__ + __rmul__ = __mul__ + + +class Force(LoadBase): + """Force acting upon a point. + + Explanation + =========== + + A force is a vector that is bound to a line of action. This class stores + both a point, which lies on the line of action, and the vector. A tuple can + also be used, with the location as the first entry and the vector as second + entry. + + Examples + ======== + + A force of magnitude 2 along N.x acting on a point Po can be created as + follows: + + >>> from sympy.physics.mechanics import Point, ReferenceFrame, Force + >>> N = ReferenceFrame('N') + >>> Po = Point('Po') + >>> Force(Po, 2 * N.x) + (Po, 2*N.x) + + If a body is supplied, then the center of mass of that body is used. + + >>> from sympy.physics.mechanics import Particle + >>> P = Particle('P', point=Po) + >>> Force(P, 2 * N.x) + (Po, 2*N.x) + + """ + + def __new__(cls, point, force): + if isinstance(point, BodyBase): + point = point.masscenter + if not isinstance(point, Point): + raise TypeError('Force location should be a Point.') + if not isinstance(force, Vector): + raise TypeError('Force vector should be a Vector.') + return super().__new__(cls, point, force) + + def __repr__(self): + return (f'{self.__class__.__name__}(point={self.point}, ' + f'force={self.force})') + + @property + def point(self): + return self.location + + @property + def force(self): + return self.vector + + +class Torque(LoadBase): + """Torque acting upon a frame. + + Explanation + =========== + + A torque is a free vector that is acting on a reference frame, which is + associated with a rigid body. This class stores both the frame and the + vector. A tuple can also be used, with the location as the first item and + the vector as second item. + + Examples + ======== + + A torque of magnitude 2 about N.x acting on a frame N can be created as + follows: + + >>> from sympy.physics.mechanics import ReferenceFrame, Torque + >>> N = ReferenceFrame('N') + >>> Torque(N, 2 * N.x) + (N, 2*N.x) + + If a body is supplied, then the frame fixed to that body is used. + + >>> from sympy.physics.mechanics import RigidBody + >>> rb = RigidBody('rb', frame=N) + >>> Torque(rb, 2 * N.x) + (N, 2*N.x) + + """ + + def __new__(cls, frame, torque): + if isinstance(frame, BodyBase): + frame = frame.frame + if not isinstance(frame, ReferenceFrame): + raise TypeError('Torque location should be a ReferenceFrame.') + if not isinstance(torque, Vector): + raise TypeError('Torque vector should be a Vector.') + return super().__new__(cls, frame, torque) + + def __repr__(self): + return (f'{self.__class__.__name__}(frame={self.frame}, ' + f'torque={self.torque})') + + @property + def frame(self): + return self.location + + @property + def torque(self): + return self.vector + + +def gravity(acceleration, *bodies): + """ + Returns a list of gravity forces given the acceleration + due to gravity and any number of particles or rigidbodies. + + Example + ======= + + >>> from sympy.physics.mechanics import ReferenceFrame, Particle, RigidBody + >>> from sympy.physics.mechanics.loads import gravity + >>> from sympy import symbols + >>> N = ReferenceFrame('N') + >>> g = symbols('g') + >>> P = Particle('P') + >>> B = RigidBody('B') + >>> gravity(g*N.y, P, B) + [(P_masscenter, P_mass*g*N.y), + (B_masscenter, B_mass*g*N.y)] + + """ + + gravity_force = [] + for body in bodies: + if not isinstance(body, BodyBase): + raise TypeError(f'{type(body)} is not a body type') + gravity_force.append(Force(body.masscenter, body.mass * acceleration)) + return gravity_force + + +def _parse_load(load): + """Helper function to parse loads and convert tuples to load objects.""" + if isinstance(load, LoadBase): + return load + elif isinstance(load, tuple): + if len(load) != 2: + raise ValueError(f'Load {load} should have a length of 2.') + if isinstance(load[0], Point): + return Force(load[0], load[1]) + elif isinstance(load[0], ReferenceFrame): + return Torque(load[0], load[1]) + else: + raise ValueError(f'Load not recognized. The load location {load[0]}' + f' should either be a Point or a ReferenceFrame.') + raise TypeError(f'Load type {type(load)} not recognized as a load. It ' + f'should be a Force, Torque or tuple.') diff --git a/lib/python3.12/site-packages/sympy/physics/mechanics/method.py b/lib/python3.12/site-packages/sympy/physics/mechanics/method.py new file mode 100644 index 0000000000000000000000000000000000000000..5c2c4a5f388e56e37bd9ecdf6daffc08ffa51070 --- /dev/null +++ b/lib/python3.12/site-packages/sympy/physics/mechanics/method.py @@ -0,0 +1,39 @@ +from abc import ABC, abstractmethod + +class _Methods(ABC): + """Abstract Base Class for all methods.""" + + @abstractmethod + def q(self): + pass + + @abstractmethod + def u(self): + pass + + @abstractmethod + def bodies(self): + pass + + @abstractmethod + def loads(self): + pass + + @abstractmethod + def mass_matrix(self): + pass + + @abstractmethod + def forcing(self): + pass + + @abstractmethod + def mass_matrix_full(self): + pass + + @abstractmethod + def forcing_full(self): + pass + + def _form_eoms(self): + raise NotImplementedError("Subclasses must implement this.") diff --git a/lib/python3.12/site-packages/sympy/physics/mechanics/pathway.py b/lib/python3.12/site-packages/sympy/physics/mechanics/pathway.py new file mode 100644 index 0000000000000000000000000000000000000000..3823750b0aeddcd8a8c4c55c02f816d823141fbb --- /dev/null +++ b/lib/python3.12/site-packages/sympy/physics/mechanics/pathway.py @@ -0,0 +1,688 @@ +"""Implementations of pathways for use by actuators.""" + +from abc import ABC, abstractmethod + +from sympy.core.singleton import S +from sympy.physics.mechanics.loads import Force +from sympy.physics.mechanics.wrapping_geometry import WrappingGeometryBase +from sympy.physics.vector import Point, dynamicsymbols + + +__all__ = ['PathwayBase', 'LinearPathway', 'ObstacleSetPathway', + 'WrappingPathway'] + + +class PathwayBase(ABC): + """Abstract base class for all pathway classes to inherit from. + + Notes + ===== + + Instances of this class cannot be directly instantiated by users. However, + it can be used to created custom pathway types through subclassing. + + """ + + def __init__(self, *attachments): + """Initializer for ``PathwayBase``.""" + self.attachments = attachments + + @property + def attachments(self): + """The pair of points defining a pathway's ends.""" + return self._attachments + + @attachments.setter + def attachments(self, attachments): + if hasattr(self, '_attachments'): + msg = ( + f'Can\'t set attribute `attachments` to {repr(attachments)} ' + f'as it is immutable.' + ) + raise AttributeError(msg) + if len(attachments) != 2: + msg = ( + f'Value {repr(attachments)} passed to `attachments` was an ' + f'iterable of length {len(attachments)}, must be an iterable ' + f'of length 2.' + ) + raise ValueError(msg) + for i, point in enumerate(attachments): + if not isinstance(point, Point): + msg = ( + f'Value {repr(point)} passed to `attachments` at index ' + f'{i} was of type {type(point)}, must be {Point}.' + ) + raise TypeError(msg) + self._attachments = tuple(attachments) + + @property + @abstractmethod + def length(self): + """An expression representing the pathway's length.""" + pass + + @property + @abstractmethod + def extension_velocity(self): + """An expression representing the pathway's extension velocity.""" + pass + + @abstractmethod + def to_loads(self, force): + """Loads required by the equations of motion method classes. + + Explanation + =========== + + ``KanesMethod`` requires a list of ``Point``-``Vector`` tuples to be + passed to the ``loads`` parameters of its ``kanes_equations`` method + when constructing the equations of motion. This method acts as a + utility to produce the correctly-structred pairs of points and vectors + required so that these can be easily concatenated with other items in + the list of loads and passed to ``KanesMethod.kanes_equations``. These + loads are also in the correct form to also be passed to the other + equations of motion method classes, e.g. ``LagrangesMethod``. + + """ + pass + + def __repr__(self): + """Default representation of a pathway.""" + attachments = ', '.join(str(a) for a in self.attachments) + return f'{self.__class__.__name__}({attachments})' + + +class LinearPathway(PathwayBase): + """Linear pathway between a pair of attachment points. + + Explanation + =========== + + A linear pathway forms a straight-line segment between two points and is + the simplest pathway that can be formed. It will not interact with any + other objects in the system, i.e. a ``LinearPathway`` will intersect other + objects to ensure that the path between its two ends (its attachments) is + the shortest possible. + + A linear pathway is made up of two points that can move relative to each + other, and a pair of equal and opposite forces acting on the points. If the + positive time-varying Euclidean distance between the two points is defined, + then the "extension velocity" is the time derivative of this distance. The + extension velocity is positive when the two points are moving away from + each other and negative when moving closer to each other. The direction for + the force acting on either point is determined by constructing a unit + vector directed from the other point to this point. This establishes a sign + convention such that a positive force magnitude tends to push the points + apart. The following diagram shows the positive force sense and the + distance between the points:: + + P Q + o<--- F --->o + | | + |<--l(t)--->| + + Examples + ======== + + >>> from sympy.physics.mechanics import LinearPathway + + To construct a pathway, two points are required to be passed to the + ``attachments`` parameter as a ``tuple``. + + >>> from sympy.physics.mechanics import Point + >>> pA, pB = Point('pA'), Point('pB') + >>> linear_pathway = LinearPathway(pA, pB) + >>> linear_pathway + LinearPathway(pA, pB) + + The pathway created above isn't very interesting without the positions and + velocities of its attachment points being described. Without this its not + possible to describe how the pathway moves, i.e. its length or its + extension velocity. + + >>> from sympy.physics.mechanics import ReferenceFrame + >>> from sympy.physics.vector import dynamicsymbols + >>> N = ReferenceFrame('N') + >>> q = dynamicsymbols('q') + >>> pB.set_pos(pA, q*N.x) + >>> pB.pos_from(pA) + q(t)*N.x + + A pathway's length can be accessed via its ``length`` attribute. + + >>> linear_pathway.length + sqrt(q(t)**2) + + Note how what appears to be an overly-complex expression is returned. This + is actually required as it ensures that a pathway's length is always + positive. + + A pathway's extension velocity can be accessed similarly via its + ``extension_velocity`` attribute. + + >>> linear_pathway.extension_velocity + sqrt(q(t)**2)*Derivative(q(t), t)/q(t) + + Parameters + ========== + + attachments : tuple[Point, Point] + Pair of ``Point`` objects between which the linear pathway spans. + Constructor expects two points to be passed, e.g. + ``LinearPathway(Point('pA'), Point('pB'))``. More or fewer points will + cause an error to be thrown. + + """ + + def __init__(self, *attachments): + """Initializer for ``LinearPathway``. + + Parameters + ========== + + attachments : Point + Pair of ``Point`` objects between which the linear pathway spans. + Constructor expects two points to be passed, e.g. + ``LinearPathway(Point('pA'), Point('pB'))``. More or fewer points + will cause an error to be thrown. + + """ + super().__init__(*attachments) + + @property + def length(self): + """Exact analytical expression for the pathway's length.""" + return _point_pair_length(*self.attachments) + + @property + def extension_velocity(self): + """Exact analytical expression for the pathway's extension velocity.""" + return _point_pair_extension_velocity(*self.attachments) + + def to_loads(self, force): + """Loads required by the equations of motion method classes. + + Explanation + =========== + + ``KanesMethod`` requires a list of ``Point``-``Vector`` tuples to be + passed to the ``loads`` parameters of its ``kanes_equations`` method + when constructing the equations of motion. This method acts as a + utility to produce the correctly-structred pairs of points and vectors + required so that these can be easily concatenated with other items in + the list of loads and passed to ``KanesMethod.kanes_equations``. These + loads are also in the correct form to also be passed to the other + equations of motion method classes, e.g. ``LagrangesMethod``. + + Examples + ======== + + The below example shows how to generate the loads produced in a linear + actuator that produces an expansile force ``F``. First, create a linear + actuator between two points separated by the coordinate ``q`` in the + ``x`` direction of the global frame ``N``. + + >>> from sympy.physics.mechanics import (LinearPathway, Point, + ... ReferenceFrame) + >>> from sympy.physics.vector import dynamicsymbols + >>> q = dynamicsymbols('q') + >>> N = ReferenceFrame('N') + >>> pA, pB = Point('pA'), Point('pB') + >>> pB.set_pos(pA, q*N.x) + >>> linear_pathway = LinearPathway(pA, pB) + + Now create a symbol ``F`` to describe the magnitude of the (expansile) + force that will be produced along the pathway. The list of loads that + ``KanesMethod`` requires can be produced by calling the pathway's + ``to_loads`` method with ``F`` passed as the only argument. + + >>> from sympy import symbols + >>> F = symbols('F') + >>> linear_pathway.to_loads(F) + [(pA, - F*q(t)/sqrt(q(t)**2)*N.x), (pB, F*q(t)/sqrt(q(t)**2)*N.x)] + + Parameters + ========== + + force : Expr + Magnitude of the force acting along the length of the pathway. As + per the sign conventions for the pathway length, pathway extension + velocity, and pair of point forces, if this ``Expr`` is positive + then the force will act to push the pair of points away from one + another (it is expansile). + + """ + relative_position = _point_pair_relative_position(*self.attachments) + loads = [ + Force(self.attachments[0], -force*relative_position/self.length), + Force(self.attachments[-1], force*relative_position/self.length), + ] + return loads + + +class ObstacleSetPathway(PathwayBase): + """Obstacle-set pathway between a set of attachment points. + + Explanation + =========== + + An obstacle-set pathway forms a series of straight-line segment between + pairs of consecutive points in a set of points. It is similiar to multiple + linear pathways joined end-to-end. It will not interact with any other + objects in the system, i.e. an ``ObstacleSetPathway`` will intersect other + objects to ensure that the path between its pairs of points (its + attachments) is the shortest possible. + + Examples + ======== + + To construct an obstacle-set pathway, three or more points are required to + be passed to the ``attachments`` parameter as a ``tuple``. + + >>> from sympy.physics.mechanics import ObstacleSetPathway, Point + >>> pA, pB, pC, pD = Point('pA'), Point('pB'), Point('pC'), Point('pD') + >>> obstacle_set_pathway = ObstacleSetPathway(pA, pB, pC, pD) + >>> obstacle_set_pathway + ObstacleSetPathway(pA, pB, pC, pD) + + The pathway created above isn't very interesting without the positions and + velocities of its attachment points being described. Without this its not + possible to describe how the pathway moves, i.e. its length or its + extension velocity. + + >>> from sympy import cos, sin + >>> from sympy.physics.mechanics import ReferenceFrame + >>> from sympy.physics.vector import dynamicsymbols + >>> N = ReferenceFrame('N') + >>> q = dynamicsymbols('q') + >>> pO = Point('pO') + >>> pA.set_pos(pO, N.y) + >>> pB.set_pos(pO, -N.x) + >>> pC.set_pos(pA, cos(q) * N.x - (sin(q) + 1) * N.y) + >>> pD.set_pos(pA, sin(q) * N.x + (cos(q) - 1) * N.y) + >>> pB.pos_from(pA) + - N.x - N.y + >>> pC.pos_from(pA) + cos(q(t))*N.x + (-sin(q(t)) - 1)*N.y + >>> pD.pos_from(pA) + sin(q(t))*N.x + (cos(q(t)) - 1)*N.y + + A pathway's length can be accessed via its ``length`` attribute. + + >>> obstacle_set_pathway.length.simplify() + sqrt(2)*(sqrt(cos(q(t)) + 1) + 2) + + A pathway's extension velocity can be accessed similarly via its + ``extension_velocity`` attribute. + + >>> obstacle_set_pathway.extension_velocity.simplify() + -sqrt(2)*sin(q(t))*Derivative(q(t), t)/(2*sqrt(cos(q(t)) + 1)) + + Parameters + ========== + + attachments : tuple[Point, Point] + The set of ``Point`` objects that define the segmented obstacle-set + pathway. + + """ + + def __init__(self, *attachments): + """Initializer for ``ObstacleSetPathway``. + + Parameters + ========== + + attachments : tuple[Point, ...] + The set of ``Point`` objects that define the segmented obstacle-set + pathway. + + """ + super().__init__(*attachments) + + @property + def attachments(self): + """The set of points defining a pathway's segmented path.""" + return self._attachments + + @attachments.setter + def attachments(self, attachments): + if hasattr(self, '_attachments'): + msg = ( + f'Can\'t set attribute `attachments` to {repr(attachments)} ' + f'as it is immutable.' + ) + raise AttributeError(msg) + if len(attachments) <= 2: + msg = ( + f'Value {repr(attachments)} passed to `attachments` was an ' + f'iterable of length {len(attachments)}, must be an iterable ' + f'of length 3 or greater.' + ) + raise ValueError(msg) + for i, point in enumerate(attachments): + if not isinstance(point, Point): + msg = ( + f'Value {repr(point)} passed to `attachments` at index ' + f'{i} was of type {type(point)}, must be {Point}.' + ) + raise TypeError(msg) + self._attachments = tuple(attachments) + + @property + def length(self): + """Exact analytical expression for the pathway's length.""" + length = S.Zero + attachment_pairs = zip(self.attachments[:-1], self.attachments[1:]) + for attachment_pair in attachment_pairs: + length += _point_pair_length(*attachment_pair) + return length + + @property + def extension_velocity(self): + """Exact analytical expression for the pathway's extension velocity.""" + extension_velocity = S.Zero + attachment_pairs = zip(self.attachments[:-1], self.attachments[1:]) + for attachment_pair in attachment_pairs: + extension_velocity += _point_pair_extension_velocity(*attachment_pair) + return extension_velocity + + def to_loads(self, force): + """Loads required by the equations of motion method classes. + + Explanation + =========== + + ``KanesMethod`` requires a list of ``Point``-``Vector`` tuples to be + passed to the ``loads`` parameters of its ``kanes_equations`` method + when constructing the equations of motion. This method acts as a + utility to produce the correctly-structred pairs of points and vectors + required so that these can be easily concatenated with other items in + the list of loads and passed to ``KanesMethod.kanes_equations``. These + loads are also in the correct form to also be passed to the other + equations of motion method classes, e.g. ``LagrangesMethod``. + + Examples + ======== + + The below example shows how to generate the loads produced in an + actuator that follows an obstacle-set pathway between four points and + produces an expansile force ``F``. First, create a pair of reference + frames, ``A`` and ``B``, in which the four points ``pA``, ``pB``, + ``pC``, and ``pD`` will be located. The first two points in frame ``A`` + and the second two in frame ``B``. Frame ``B`` will also be oriented + such that it relates to ``A`` via a rotation of ``q`` about an axis + ``N.z`` in a global frame (``N.z``, ``A.z``, and ``B.z`` are parallel). + + >>> from sympy.physics.mechanics import (ObstacleSetPathway, Point, + ... ReferenceFrame) + >>> from sympy.physics.vector import dynamicsymbols + >>> q = dynamicsymbols('q') + >>> N = ReferenceFrame('N') + >>> N = ReferenceFrame('N') + >>> A = N.orientnew('A', 'axis', (0, N.x)) + >>> B = A.orientnew('B', 'axis', (q, N.z)) + >>> pO = Point('pO') + >>> pA, pB, pC, pD = Point('pA'), Point('pB'), Point('pC'), Point('pD') + >>> pA.set_pos(pO, A.x) + >>> pB.set_pos(pO, -A.y) + >>> pC.set_pos(pO, B.y) + >>> pD.set_pos(pO, B.x) + >>> obstacle_set_pathway = ObstacleSetPathway(pA, pB, pC, pD) + + Now create a symbol ``F`` to describe the magnitude of the (expansile) + force that will be produced along the pathway. The list of loads that + ``KanesMethod`` requires can be produced by calling the pathway's + ``to_loads`` method with ``F`` passed as the only argument. + + >>> from sympy import Symbol + >>> F = Symbol('F') + >>> obstacle_set_pathway.to_loads(F) + [(pA, sqrt(2)*F/2*A.x + sqrt(2)*F/2*A.y), + (pB, - sqrt(2)*F/2*A.x - sqrt(2)*F/2*A.y), + (pB, - F/sqrt(2*cos(q(t)) + 2)*A.y - F/sqrt(2*cos(q(t)) + 2)*B.y), + (pC, F/sqrt(2*cos(q(t)) + 2)*A.y + F/sqrt(2*cos(q(t)) + 2)*B.y), + (pC, - sqrt(2)*F/2*B.x + sqrt(2)*F/2*B.y), + (pD, sqrt(2)*F/2*B.x - sqrt(2)*F/2*B.y)] + + Parameters + ========== + + force : Expr + The force acting along the length of the pathway. It is assumed + that this ``Expr`` represents an expansile force. + + """ + loads = [] + attachment_pairs = zip(self.attachments[:-1], self.attachments[1:]) + for attachment_pair in attachment_pairs: + relative_position = _point_pair_relative_position(*attachment_pair) + length = _point_pair_length(*attachment_pair) + loads.extend([ + Force(attachment_pair[0], -force*relative_position/length), + Force(attachment_pair[1], force*relative_position/length), + ]) + return loads + + +class WrappingPathway(PathwayBase): + """Pathway that wraps a geometry object. + + Explanation + =========== + + A wrapping pathway interacts with a geometry object and forms a path that + wraps smoothly along its surface. The wrapping pathway along the geometry + object will be the geodesic that the geometry object defines based on the + two points. It will not interact with any other objects in the system, i.e. + a ``WrappingPathway`` will intersect other objects to ensure that the path + between its two ends (its attachments) is the shortest possible. + + To explain the sign conventions used for pathway length, extension + velocity, and direction of applied forces, we can ignore the geometry with + which the wrapping pathway interacts. A wrapping pathway is made up of two + points that can move relative to each other, and a pair of equal and + opposite forces acting on the points. If the positive time-varying + Euclidean distance between the two points is defined, then the "extension + velocity" is the time derivative of this distance. The extension velocity + is positive when the two points are moving away from each other and + negative when moving closer to each other. The direction for the force + acting on either point is determined by constructing a unit vector directed + from the other point to this point. This establishes a sign convention such + that a positive force magnitude tends to push the points apart. The + following diagram shows the positive force sense and the distance between + the points:: + + P Q + o<--- F --->o + | | + |<--l(t)--->| + + Examples + ======== + + >>> from sympy.physics.mechanics import WrappingPathway + + To construct a wrapping pathway, like other pathways, a pair of points must + be passed, followed by an instance of a wrapping geometry class as a + keyword argument. We'll use a cylinder with radius ``r`` and its axis + parallel to ``N.x`` passing through a point ``pO``. + + >>> from sympy import symbols + >>> from sympy.physics.mechanics import Point, ReferenceFrame, WrappingCylinder + >>> r = symbols('r') + >>> N = ReferenceFrame('N') + >>> pA, pB, pO = Point('pA'), Point('pB'), Point('pO') + >>> cylinder = WrappingCylinder(r, pO, N.x) + >>> wrapping_pathway = WrappingPathway(pA, pB, cylinder) + >>> wrapping_pathway + WrappingPathway(pA, pB, geometry=WrappingCylinder(radius=r, point=pO, + axis=N.x)) + + Parameters + ========== + + attachment_1 : Point + First of the pair of ``Point`` objects between which the wrapping + pathway spans. + attachment_2 : Point + Second of the pair of ``Point`` objects between which the wrapping + pathway spans. + geometry : WrappingGeometryBase + Geometry about which the pathway wraps. + + """ + + def __init__(self, attachment_1, attachment_2, geometry): + """Initializer for ``WrappingPathway``. + + Parameters + ========== + + attachment_1 : Point + First of the pair of ``Point`` objects between which the wrapping + pathway spans. + attachment_2 : Point + Second of the pair of ``Point`` objects between which the wrapping + pathway spans. + geometry : WrappingGeometryBase + Geometry about which the pathway wraps. + The geometry about which the pathway wraps. + + """ + super().__init__(attachment_1, attachment_2) + self.geometry = geometry + + @property + def geometry(self): + """Geometry around which the pathway wraps.""" + return self._geometry + + @geometry.setter + def geometry(self, geometry): + if hasattr(self, '_geometry'): + msg = ( + f'Can\'t set attribute `geometry` to {repr(geometry)} as it ' + f'is immutable.' + ) + raise AttributeError(msg) + if not isinstance(geometry, WrappingGeometryBase): + msg = ( + f'Value {repr(geometry)} passed to `geometry` was of type ' + f'{type(geometry)}, must be {WrappingGeometryBase}.' + ) + raise TypeError(msg) + self._geometry = geometry + + @property + def length(self): + """Exact analytical expression for the pathway's length.""" + return self.geometry.geodesic_length(*self.attachments) + + @property + def extension_velocity(self): + """Exact analytical expression for the pathway's extension velocity.""" + return self.length.diff(dynamicsymbols._t) + + def to_loads(self, force): + """Loads required by the equations of motion method classes. + + Explanation + =========== + + ``KanesMethod`` requires a list of ``Point``-``Vector`` tuples to be + passed to the ``loads`` parameters of its ``kanes_equations`` method + when constructing the equations of motion. This method acts as a + utility to produce the correctly-structred pairs of points and vectors + required so that these can be easily concatenated with other items in + the list of loads and passed to ``KanesMethod.kanes_equations``. These + loads are also in the correct form to also be passed to the other + equations of motion method classes, e.g. ``LagrangesMethod``. + + Examples + ======== + + The below example shows how to generate the loads produced in an + actuator that produces an expansile force ``F`` while wrapping around a + cylinder. First, create a cylinder with radius ``r`` and an axis + parallel to the ``N.z`` direction of the global frame ``N`` that also + passes through a point ``pO``. + + >>> from sympy import symbols + >>> from sympy.physics.mechanics import (Point, ReferenceFrame, + ... WrappingCylinder) + >>> N = ReferenceFrame('N') + >>> r = symbols('r', positive=True) + >>> pO = Point('pO') + >>> cylinder = WrappingCylinder(r, pO, N.z) + + Create the pathway of the actuator using the ``WrappingPathway`` class, + defined to span between two points ``pA`` and ``pB``. Both points lie + on the surface of the cylinder and the location of ``pB`` is defined + relative to ``pA`` by the dynamics symbol ``q``. + + >>> from sympy import cos, sin + >>> from sympy.physics.mechanics import WrappingPathway, dynamicsymbols + >>> q = dynamicsymbols('q') + >>> pA = Point('pA') + >>> pB = Point('pB') + >>> pA.set_pos(pO, r*N.x) + >>> pB.set_pos(pO, r*(cos(q)*N.x + sin(q)*N.y)) + >>> pB.pos_from(pA) + (r*cos(q(t)) - r)*N.x + r*sin(q(t))*N.y + >>> pathway = WrappingPathway(pA, pB, cylinder) + + Now create a symbol ``F`` to describe the magnitude of the (expansile) + force that will be produced along the pathway. The list of loads that + ``KanesMethod`` requires can be produced by calling the pathway's + ``to_loads`` method with ``F`` passed as the only argument. + + >>> F = symbols('F') + >>> loads = pathway.to_loads(F) + >>> [load.__class__(load.location, load.vector.simplify()) for load in loads] + [(pA, F*N.y), (pB, F*sin(q(t))*N.x - F*cos(q(t))*N.y), + (pO, - F*sin(q(t))*N.x + F*(cos(q(t)) - 1)*N.y)] + + Parameters + ========== + + force : Expr + Magnitude of the force acting along the length of the pathway. It + is assumed that this ``Expr`` represents an expansile force. + + """ + pA, pB = self.attachments + pO = self.geometry.point + pA_force, pB_force = self.geometry.geodesic_end_vectors(pA, pB) + pO_force = -(pA_force + pB_force) + + loads = [ + Force(pA, force * pA_force), + Force(pB, force * pB_force), + Force(pO, force * pO_force), + ] + return loads + + def __repr__(self): + """Representation of a ``WrappingPathway``.""" + attachments = ', '.join(str(a) for a in self.attachments) + return ( + f'{self.__class__.__name__}({attachments}, ' + f'geometry={self.geometry})' + ) + + +def _point_pair_relative_position(point_1, point_2): + """The relative position between a pair of points.""" + return point_2.pos_from(point_1) + + +def _point_pair_length(point_1, point_2): + """The length of the direct linear path between two points.""" + return _point_pair_relative_position(point_1, point_2).magnitude() + + +def _point_pair_extension_velocity(point_1, point_2): + """The extension velocity of the direct linear path between two points.""" + return _point_pair_length(point_1, point_2).diff(dynamicsymbols._t) diff --git a/lib/python3.12/site-packages/sympy/physics/mechanics/tests/__init__.py b/lib/python3.12/site-packages/sympy/physics/mechanics/tests/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/sympy/physics/mechanics/tests/__pycache__/test_inertia.cpython-312.pyc b/lib/python3.12/site-packages/sympy/physics/mechanics/tests/__pycache__/test_inertia.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6982af3531c79d242e00f0912731e9770529dc84 Binary files /dev/null and b/lib/python3.12/site-packages/sympy/physics/mechanics/tests/__pycache__/test_inertia.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sympy/physics/mechanics/tests/__pycache__/test_kane5.cpython-312.pyc b/lib/python3.12/site-packages/sympy/physics/mechanics/tests/__pycache__/test_kane5.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..eea84b9577d64de83d5d17b1ca0890b795c0857a Binary files /dev/null and b/lib/python3.12/site-packages/sympy/physics/mechanics/tests/__pycache__/test_kane5.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sympy/physics/mechanics/tests/__pycache__/test_method.cpython-312.pyc b/lib/python3.12/site-packages/sympy/physics/mechanics/tests/__pycache__/test_method.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..284e869e231c8fd7c247dbc6d4c47b1ba588b7dd Binary files /dev/null and b/lib/python3.12/site-packages/sympy/physics/mechanics/tests/__pycache__/test_method.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sympy/physics/mechanics/tests/__pycache__/test_rigidbody.cpython-312.pyc b/lib/python3.12/site-packages/sympy/physics/mechanics/tests/__pycache__/test_rigidbody.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..70faedc1599adaf60bc6b01c36c6fbf0d622506c Binary files /dev/null and b/lib/python3.12/site-packages/sympy/physics/mechanics/tests/__pycache__/test_rigidbody.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sympy/physics/mechanics/tests/__pycache__/test_system.cpython-312.pyc b/lib/python3.12/site-packages/sympy/physics/mechanics/tests/__pycache__/test_system.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9d0b0dcd0e22abe37b92b5d484340c1a5b082f0e Binary files /dev/null and b/lib/python3.12/site-packages/sympy/physics/mechanics/tests/__pycache__/test_system.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sympy/physics/mechanics/tests/test_actuator.py b/lib/python3.12/site-packages/sympy/physics/mechanics/tests/test_actuator.py new file mode 100644 index 0000000000000000000000000000000000000000..9fc9fa21164e98ebb452aabe04ec27d76595e96c --- /dev/null +++ b/lib/python3.12/site-packages/sympy/physics/mechanics/tests/test_actuator.py @@ -0,0 +1,851 @@ +"""Tests for the ``sympy.physics.mechanics.actuator.py`` module.""" + +import pytest + +from sympy import ( + S, + Matrix, + Symbol, + SympifyError, + sqrt, + Abs +) +from sympy.physics.mechanics import ( + ActuatorBase, + Force, + ForceActuator, + KanesMethod, + LinearDamper, + LinearPathway, + LinearSpring, + Particle, + PinJoint, + Point, + ReferenceFrame, + RigidBody, + TorqueActuator, + Vector, + dynamicsymbols, + DuffingSpring, +) + +from sympy.core.expr import Expr as ExprType + +target = RigidBody('target') +reaction = RigidBody('reaction') + + +class TestForceActuator: + + @pytest.fixture(autouse=True) + def _linear_pathway_fixture(self): + self.force = Symbol('F') + self.pA = Point('pA') + self.pB = Point('pB') + self.pathway = LinearPathway(self.pA, self.pB) + self.q1 = dynamicsymbols('q1') + self.q2 = dynamicsymbols('q2') + self.q3 = dynamicsymbols('q3') + self.q1d = dynamicsymbols('q1', 1) + self.q2d = dynamicsymbols('q2', 1) + self.q3d = dynamicsymbols('q3', 1) + self.N = ReferenceFrame('N') + + def test_is_actuator_base_subclass(self): + assert issubclass(ForceActuator, ActuatorBase) + + @pytest.mark.parametrize( + 'force, expected_force', + [ + (1, S.One), + (S.One, S.One), + (Symbol('F'), Symbol('F')), + (dynamicsymbols('F'), dynamicsymbols('F')), + (Symbol('F')**2 + Symbol('F'), Symbol('F')**2 + Symbol('F')), + ] + ) + def test_valid_constructor_force(self, force, expected_force): + instance = ForceActuator(force, self.pathway) + assert isinstance(instance, ForceActuator) + assert hasattr(instance, 'force') + assert isinstance(instance.force, ExprType) + assert instance.force == expected_force + + @pytest.mark.parametrize('force', [None, 'F']) + def test_invalid_constructor_force_not_sympifyable(self, force): + with pytest.raises(SympifyError): + _ = ForceActuator(force, self.pathway) + + @pytest.mark.parametrize( + 'pathway', + [ + LinearPathway(Point('pA'), Point('pB')), + ] + ) + def test_valid_constructor_pathway(self, pathway): + instance = ForceActuator(self.force, pathway) + assert isinstance(instance, ForceActuator) + assert hasattr(instance, 'pathway') + assert isinstance(instance.pathway, LinearPathway) + assert instance.pathway == pathway + + def test_invalid_constructor_pathway_not_pathway_base(self): + with pytest.raises(TypeError): + _ = ForceActuator(self.force, None) + + @pytest.mark.parametrize( + 'property_name, fixture_attr_name', + [ + ('force', 'force'), + ('pathway', 'pathway'), + ] + ) + def test_properties_are_immutable(self, property_name, fixture_attr_name): + instance = ForceActuator(self.force, self.pathway) + value = getattr(self, fixture_attr_name) + with pytest.raises(AttributeError): + setattr(instance, property_name, value) + + def test_repr(self): + actuator = ForceActuator(self.force, self.pathway) + expected = "ForceActuator(F, LinearPathway(pA, pB))" + assert repr(actuator) == expected + + def test_to_loads_static_pathway(self): + self.pB.set_pos(self.pA, 2*self.N.x) + actuator = ForceActuator(self.force, self.pathway) + expected = [ + (self.pA, - self.force*self.N.x), + (self.pB, self.force*self.N.x), + ] + assert actuator.to_loads() == expected + + def test_to_loads_2D_pathway(self): + self.pB.set_pos(self.pA, 2*self.q1*self.N.x) + actuator = ForceActuator(self.force, self.pathway) + expected = [ + (self.pA, - self.force*(self.q1/sqrt(self.q1**2))*self.N.x), + (self.pB, self.force*(self.q1/sqrt(self.q1**2))*self.N.x), + ] + assert actuator.to_loads() == expected + + def test_to_loads_3D_pathway(self): + self.pB.set_pos( + self.pA, + self.q1*self.N.x - self.q2*self.N.y + 2*self.q3*self.N.z, + ) + actuator = ForceActuator(self.force, self.pathway) + length = sqrt(self.q1**2 + self.q2**2 + 4*self.q3**2) + pO_force = ( + - self.force*self.q1*self.N.x/length + + self.force*self.q2*self.N.y/length + - 2*self.force*self.q3*self.N.z/length + ) + pI_force = ( + self.force*self.q1*self.N.x/length + - self.force*self.q2*self.N.y/length + + 2*self.force*self.q3*self.N.z/length + ) + expected = [ + (self.pA, pO_force), + (self.pB, pI_force), + ] + assert actuator.to_loads() == expected + + +class TestLinearSpring: + + @pytest.fixture(autouse=True) + def _linear_spring_fixture(self): + self.stiffness = Symbol('k') + self.l = Symbol('l') + self.pA = Point('pA') + self.pB = Point('pB') + self.pathway = LinearPathway(self.pA, self.pB) + self.q = dynamicsymbols('q') + self.N = ReferenceFrame('N') + + def test_is_force_actuator_subclass(self): + assert issubclass(LinearSpring, ForceActuator) + + def test_is_actuator_base_subclass(self): + assert issubclass(LinearSpring, ActuatorBase) + + @pytest.mark.parametrize( + ( + 'stiffness, ' + 'expected_stiffness, ' + 'equilibrium_length, ' + 'expected_equilibrium_length, ' + 'force' + ), + [ + ( + 1, + S.One, + 0, + S.Zero, + -sqrt(dynamicsymbols('q')**2), + ), + ( + Symbol('k'), + Symbol('k'), + 0, + S.Zero, + -Symbol('k')*sqrt(dynamicsymbols('q')**2), + ), + ( + Symbol('k'), + Symbol('k'), + S.Zero, + S.Zero, + -Symbol('k')*sqrt(dynamicsymbols('q')**2), + ), + ( + Symbol('k'), + Symbol('k'), + Symbol('l'), + Symbol('l'), + -Symbol('k')*(sqrt(dynamicsymbols('q')**2) - Symbol('l')), + ), + ] + ) + def test_valid_constructor( + self, + stiffness, + expected_stiffness, + equilibrium_length, + expected_equilibrium_length, + force, + ): + self.pB.set_pos(self.pA, self.q*self.N.x) + spring = LinearSpring(stiffness, self.pathway, equilibrium_length) + + assert isinstance(spring, LinearSpring) + + assert hasattr(spring, 'stiffness') + assert isinstance(spring.stiffness, ExprType) + assert spring.stiffness == expected_stiffness + + assert hasattr(spring, 'pathway') + assert isinstance(spring.pathway, LinearPathway) + assert spring.pathway == self.pathway + + assert hasattr(spring, 'equilibrium_length') + assert isinstance(spring.equilibrium_length, ExprType) + assert spring.equilibrium_length == expected_equilibrium_length + + assert hasattr(spring, 'force') + assert isinstance(spring.force, ExprType) + assert spring.force == force + + @pytest.mark.parametrize('stiffness', [None, 'k']) + def test_invalid_constructor_stiffness_not_sympifyable(self, stiffness): + with pytest.raises(SympifyError): + _ = LinearSpring(stiffness, self.pathway, self.l) + + def test_invalid_constructor_pathway_not_pathway_base(self): + with pytest.raises(TypeError): + _ = LinearSpring(self.stiffness, None, self.l) + + @pytest.mark.parametrize('equilibrium_length', [None, 'l']) + def test_invalid_constructor_equilibrium_length_not_sympifyable( + self, + equilibrium_length, + ): + with pytest.raises(SympifyError): + _ = LinearSpring(self.stiffness, self.pathway, equilibrium_length) + + @pytest.mark.parametrize( + 'property_name, fixture_attr_name', + [ + ('stiffness', 'stiffness'), + ('pathway', 'pathway'), + ('equilibrium_length', 'l'), + ] + ) + def test_properties_are_immutable(self, property_name, fixture_attr_name): + spring = LinearSpring(self.stiffness, self.pathway, self.l) + value = getattr(self, fixture_attr_name) + with pytest.raises(AttributeError): + setattr(spring, property_name, value) + + @pytest.mark.parametrize( + 'equilibrium_length, expected', + [ + (S.Zero, 'LinearSpring(k, LinearPathway(pA, pB))'), + ( + Symbol('l'), + 'LinearSpring(k, LinearPathway(pA, pB), equilibrium_length=l)', + ), + ] + ) + def test_repr(self, equilibrium_length, expected): + self.pB.set_pos(self.pA, self.q*self.N.x) + spring = LinearSpring(self.stiffness, self.pathway, equilibrium_length) + assert repr(spring) == expected + + def test_to_loads(self): + self.pB.set_pos(self.pA, self.q*self.N.x) + spring = LinearSpring(self.stiffness, self.pathway, self.l) + normal = self.q/sqrt(self.q**2)*self.N.x + pA_force = self.stiffness*(sqrt(self.q**2) - self.l)*normal + pB_force = -self.stiffness*(sqrt(self.q**2) - self.l)*normal + expected = [Force(self.pA, pA_force), Force(self.pB, pB_force)] + loads = spring.to_loads() + + for load, (point, vector) in zip(loads, expected): + assert isinstance(load, Force) + assert load.point == point + assert (load.vector - vector).simplify() == 0 + + +class TestLinearDamper: + + @pytest.fixture(autouse=True) + def _linear_damper_fixture(self): + self.damping = Symbol('c') + self.l = Symbol('l') + self.pA = Point('pA') + self.pB = Point('pB') + self.pathway = LinearPathway(self.pA, self.pB) + self.q = dynamicsymbols('q') + self.dq = dynamicsymbols('q', 1) + self.u = dynamicsymbols('u') + self.N = ReferenceFrame('N') + + def test_is_force_actuator_subclass(self): + assert issubclass(LinearDamper, ForceActuator) + + def test_is_actuator_base_subclass(self): + assert issubclass(LinearDamper, ActuatorBase) + + def test_valid_constructor(self): + self.pB.set_pos(self.pA, self.q*self.N.x) + damper = LinearDamper(self.damping, self.pathway) + + assert isinstance(damper, LinearDamper) + + assert hasattr(damper, 'damping') + assert isinstance(damper.damping, ExprType) + assert damper.damping == self.damping + + assert hasattr(damper, 'pathway') + assert isinstance(damper.pathway, LinearPathway) + assert damper.pathway == self.pathway + + def test_valid_constructor_force(self): + self.pB.set_pos(self.pA, self.q*self.N.x) + damper = LinearDamper(self.damping, self.pathway) + + expected_force = -self.damping*sqrt(self.q**2)*self.dq/self.q + assert hasattr(damper, 'force') + assert isinstance(damper.force, ExprType) + assert damper.force == expected_force + + @pytest.mark.parametrize('damping', [None, 'c']) + def test_invalid_constructor_damping_not_sympifyable(self, damping): + with pytest.raises(SympifyError): + _ = LinearDamper(damping, self.pathway) + + def test_invalid_constructor_pathway_not_pathway_base(self): + with pytest.raises(TypeError): + _ = LinearDamper(self.damping, None) + + @pytest.mark.parametrize( + 'property_name, fixture_attr_name', + [ + ('damping', 'damping'), + ('pathway', 'pathway'), + ] + ) + def test_properties_are_immutable(self, property_name, fixture_attr_name): + damper = LinearDamper(self.damping, self.pathway) + value = getattr(self, fixture_attr_name) + with pytest.raises(AttributeError): + setattr(damper, property_name, value) + + def test_repr(self): + self.pB.set_pos(self.pA, self.q*self.N.x) + damper = LinearDamper(self.damping, self.pathway) + expected = 'LinearDamper(c, LinearPathway(pA, pB))' + assert repr(damper) == expected + + def test_to_loads(self): + self.pB.set_pos(self.pA, self.q*self.N.x) + damper = LinearDamper(self.damping, self.pathway) + direction = self.q**2/self.q**2*self.N.x + pA_force = self.damping*self.dq*direction + pB_force = -self.damping*self.dq*direction + expected = [Force(self.pA, pA_force), Force(self.pB, pB_force)] + assert damper.to_loads() == expected + + +class TestForcedMassSpringDamperModel(): + r"""A single degree of freedom translational forced mass-spring-damper. + + Notes + ===== + + This system is well known to have the governing equation: + + .. math:: + m \ddot{x} = F - k x - c \dot{x} + + where $F$ is an externally applied force, $m$ is the mass of the particle + to which the spring and damper are attached, $k$ is the spring's stiffness, + $c$ is the dampers damping coefficient, and $x$ is the generalized + coordinate representing the system's single (translational) degree of + freedom. + + """ + + @pytest.fixture(autouse=True) + def _force_mass_spring_damper_model_fixture(self): + self.m = Symbol('m') + self.k = Symbol('k') + self.c = Symbol('c') + self.F = Symbol('F') + + self.q = dynamicsymbols('q') + self.dq = dynamicsymbols('q', 1) + self.u = dynamicsymbols('u') + + self.frame = ReferenceFrame('N') + self.origin = Point('pO') + self.origin.set_vel(self.frame, 0) + + self.attachment = Point('pA') + self.attachment.set_pos(self.origin, self.q*self.frame.x) + + self.mass = Particle('mass', self.attachment, self.m) + self.pathway = LinearPathway(self.origin, self.attachment) + + self.kanes_method = KanesMethod( + self.frame, + q_ind=[self.q], + u_ind=[self.u], + kd_eqs=[self.dq - self.u], + ) + self.bodies = [self.mass] + + self.mass_matrix = Matrix([[self.m]]) + self.forcing = Matrix([[self.F - self.c*self.u - self.k*self.q]]) + + def test_force_acuator(self): + stiffness = -self.k*self.pathway.length + spring = ForceActuator(stiffness, self.pathway) + damping = -self.c*self.pathway.extension_velocity + damper = ForceActuator(damping, self.pathway) + + loads = [ + (self.attachment, self.F*self.frame.x), + *spring.to_loads(), + *damper.to_loads(), + ] + self.kanes_method.kanes_equations(self.bodies, loads) + + assert self.kanes_method.mass_matrix == self.mass_matrix + assert self.kanes_method.forcing == self.forcing + + def test_linear_spring_linear_damper(self): + spring = LinearSpring(self.k, self.pathway) + damper = LinearDamper(self.c, self.pathway) + + loads = [ + (self.attachment, self.F*self.frame.x), + *spring.to_loads(), + *damper.to_loads(), + ] + self.kanes_method.kanes_equations(self.bodies, loads) + + assert self.kanes_method.mass_matrix == self.mass_matrix + assert self.kanes_method.forcing == self.forcing + + +class TestTorqueActuator: + + @pytest.fixture(autouse=True) + def _torque_actuator_fixture(self): + self.torque = Symbol('T') + self.N = ReferenceFrame('N') + self.A = ReferenceFrame('A') + self.axis = self.N.z + self.target = RigidBody('target', frame=self.N) + self.reaction = RigidBody('reaction', frame=self.A) + + def test_is_actuator_base_subclass(self): + assert issubclass(TorqueActuator, ActuatorBase) + + @pytest.mark.parametrize( + 'torque', + [ + Symbol('T'), + dynamicsymbols('T'), + Symbol('T')**2 + Symbol('T'), + ] + ) + @pytest.mark.parametrize( + 'target_frame, reaction_frame', + [ + (target.frame, reaction.frame), + (target, reaction.frame), + (target.frame, reaction), + (target, reaction), + ] + ) + def test_valid_constructor_with_reaction( + self, + torque, + target_frame, + reaction_frame, + ): + instance = TorqueActuator( + torque, + self.axis, + target_frame, + reaction_frame, + ) + assert isinstance(instance, TorqueActuator) + + assert hasattr(instance, 'torque') + assert isinstance(instance.torque, ExprType) + assert instance.torque == torque + + assert hasattr(instance, 'axis') + assert isinstance(instance.axis, Vector) + assert instance.axis == self.axis + + assert hasattr(instance, 'target_frame') + assert isinstance(instance.target_frame, ReferenceFrame) + assert instance.target_frame == target.frame + + assert hasattr(instance, 'reaction_frame') + assert isinstance(instance.reaction_frame, ReferenceFrame) + assert instance.reaction_frame == reaction.frame + + @pytest.mark.parametrize( + 'torque', + [ + Symbol('T'), + dynamicsymbols('T'), + Symbol('T')**2 + Symbol('T'), + ] + ) + @pytest.mark.parametrize('target_frame', [target.frame, target]) + def test_valid_constructor_without_reaction(self, torque, target_frame): + instance = TorqueActuator(torque, self.axis, target_frame) + assert isinstance(instance, TorqueActuator) + + assert hasattr(instance, 'torque') + assert isinstance(instance.torque, ExprType) + assert instance.torque == torque + + assert hasattr(instance, 'axis') + assert isinstance(instance.axis, Vector) + assert instance.axis == self.axis + + assert hasattr(instance, 'target_frame') + assert isinstance(instance.target_frame, ReferenceFrame) + assert instance.target_frame == target.frame + + assert hasattr(instance, 'reaction_frame') + assert instance.reaction_frame is None + + @pytest.mark.parametrize('torque', [None, 'T']) + def test_invalid_constructor_torque_not_sympifyable(self, torque): + with pytest.raises(SympifyError): + _ = TorqueActuator(torque, self.axis, self.target) + + @pytest.mark.parametrize('axis', [Symbol('a'), dynamicsymbols('a')]) + def test_invalid_constructor_axis_not_vector(self, axis): + with pytest.raises(TypeError): + _ = TorqueActuator(self.torque, axis, self.target, self.reaction) + + @pytest.mark.parametrize( + 'frames', + [ + (None, ReferenceFrame('child')), + (ReferenceFrame('parent'), True), + (None, RigidBody('child')), + (RigidBody('parent'), True), + ] + ) + def test_invalid_constructor_frames_not_frame(self, frames): + with pytest.raises(TypeError): + _ = TorqueActuator(self.torque, self.axis, *frames) + + @pytest.mark.parametrize( + 'property_name, fixture_attr_name', + [ + ('torque', 'torque'), + ('axis', 'axis'), + ('target_frame', 'target'), + ('reaction_frame', 'reaction'), + ] + ) + def test_properties_are_immutable(self, property_name, fixture_attr_name): + actuator = TorqueActuator( + self.torque, + self.axis, + self.target, + self.reaction, + ) + value = getattr(self, fixture_attr_name) + with pytest.raises(AttributeError): + setattr(actuator, property_name, value) + + def test_repr_without_reaction(self): + actuator = TorqueActuator(self.torque, self.axis, self.target) + expected = 'TorqueActuator(T, axis=N.z, target_frame=N)' + assert repr(actuator) == expected + + def test_repr_with_reaction(self): + actuator = TorqueActuator( + self.torque, + self.axis, + self.target, + self.reaction, + ) + expected = 'TorqueActuator(T, axis=N.z, target_frame=N, reaction_frame=A)' + assert repr(actuator) == expected + + def test_at_pin_joint_constructor(self): + pin_joint = PinJoint( + 'pin', + self.target, + self.reaction, + coordinates=dynamicsymbols('q'), + speeds=dynamicsymbols('u'), + parent_interframe=self.N, + joint_axis=self.axis, + ) + instance = TorqueActuator.at_pin_joint(self.torque, pin_joint) + assert isinstance(instance, TorqueActuator) + + assert hasattr(instance, 'torque') + assert isinstance(instance.torque, ExprType) + assert instance.torque == self.torque + + assert hasattr(instance, 'axis') + assert isinstance(instance.axis, Vector) + assert instance.axis == self.axis + + assert hasattr(instance, 'target_frame') + assert isinstance(instance.target_frame, ReferenceFrame) + assert instance.target_frame == self.A + + assert hasattr(instance, 'reaction_frame') + assert isinstance(instance.reaction_frame, ReferenceFrame) + assert instance.reaction_frame == self.N + + def test_at_pin_joint_pin_joint_not_pin_joint_invalid(self): + with pytest.raises(TypeError): + _ = TorqueActuator.at_pin_joint(self.torque, Symbol('pin')) + + def test_to_loads_without_reaction(self): + actuator = TorqueActuator(self.torque, self.axis, self.target) + expected = [ + (self.N, self.torque*self.axis), + ] + assert actuator.to_loads() == expected + + def test_to_loads_with_reaction(self): + actuator = TorqueActuator( + self.torque, + self.axis, + self.target, + self.reaction, + ) + expected = [ + (self.N, self.torque*self.axis), + (self.A, - self.torque*self.axis), + ] + assert actuator.to_loads() == expected + + +class NonSympifyable: + pass + + +class TestDuffingSpring: + @pytest.fixture(autouse=True) + # Set up common vairables that will be used in multiple tests + def _duffing_spring_fixture(self): + self.linear_stiffness = Symbol('beta') + self.nonlinear_stiffness = Symbol('alpha') + self.equilibrium_length = Symbol('l') + self.pA = Point('pA') + self.pB = Point('pB') + self.pathway = LinearPathway(self.pA, self.pB) + self.q = dynamicsymbols('q') + self.N = ReferenceFrame('N') + + # Simples tests to check that DuffingSpring is a subclass of ForceActuator and ActuatorBase + def test_is_force_actuator_subclass(self): + assert issubclass(DuffingSpring, ForceActuator) + + def test_is_actuator_base_subclass(self): + assert issubclass(DuffingSpring, ActuatorBase) + + @pytest.mark.parametrize( + # Create parametrized tests that allows running the same test function multiple times with different sets of arguments + ( + 'linear_stiffness, ' + 'expected_linear_stiffness, ' + 'nonlinear_stiffness, ' + 'expected_nonlinear_stiffness, ' + 'equilibrium_length, ' + 'expected_equilibrium_length, ' + 'force' + ), + [ + ( + 1, + S.One, + 1, + S.One, + 0, + S.Zero, + -sqrt(dynamicsymbols('q')**2)-(sqrt(dynamicsymbols('q')**2))**3, + ), + ( + Symbol('beta'), + Symbol('beta'), + Symbol('alpha'), + Symbol('alpha'), + 0, + S.Zero, + -Symbol('beta')*sqrt(dynamicsymbols('q')**2)-Symbol('alpha')*(sqrt(dynamicsymbols('q')**2))**3, + ), + ( + Symbol('beta'), + Symbol('beta'), + Symbol('alpha'), + Symbol('alpha'), + S.Zero, + S.Zero, + -Symbol('beta')*sqrt(dynamicsymbols('q')**2)-Symbol('alpha')*(sqrt(dynamicsymbols('q')**2))**3, + ), + ( + Symbol('beta'), + Symbol('beta'), + Symbol('alpha'), + Symbol('alpha'), + Symbol('l'), + Symbol('l'), + -Symbol('beta') * (sqrt(dynamicsymbols('q')**2) - Symbol('l')) - Symbol('alpha') * (sqrt(dynamicsymbols('q')**2) - Symbol('l'))**3, + ), + ] + ) + + # Check if DuffingSpring correctly inializes its attributes + # It tests various combinations of linear & nonlinear stiffness, equilibriun length, and the resulting force expression + def test_valid_constructor( + self, + linear_stiffness, + expected_linear_stiffness, + nonlinear_stiffness, + expected_nonlinear_stiffness, + equilibrium_length, + expected_equilibrium_length, + force, + ): + self.pB.set_pos(self.pA, self.q*self.N.x) + spring = DuffingSpring(linear_stiffness, nonlinear_stiffness, self.pathway, equilibrium_length) + + assert isinstance(spring, DuffingSpring) + + assert hasattr(spring, 'linear_stiffness') + assert isinstance(spring.linear_stiffness, ExprType) + assert spring.linear_stiffness == expected_linear_stiffness + + assert hasattr(spring, 'nonlinear_stiffness') + assert isinstance(spring.nonlinear_stiffness, ExprType) + assert spring.nonlinear_stiffness == expected_nonlinear_stiffness + + assert hasattr(spring, 'pathway') + assert isinstance(spring.pathway, LinearPathway) + assert spring.pathway == self.pathway + + assert hasattr(spring, 'equilibrium_length') + assert isinstance(spring.equilibrium_length, ExprType) + assert spring.equilibrium_length == expected_equilibrium_length + + assert hasattr(spring, 'force') + assert isinstance(spring.force, ExprType) + assert spring.force == force + + @pytest.mark.parametrize('linear_stiffness', [None, NonSympifyable()]) + def test_invalid_constructor_linear_stiffness_not_sympifyable(self, linear_stiffness): + with pytest.raises(SympifyError): + _ = DuffingSpring(linear_stiffness, self.nonlinear_stiffness, self.pathway, self.equilibrium_length) + + @pytest.mark.parametrize('nonlinear_stiffness', [None, NonSympifyable()]) + def test_invalid_constructor_nonlinear_stiffness_not_sympifyable(self, nonlinear_stiffness): + with pytest.raises(SympifyError): + _ = DuffingSpring(self.linear_stiffness, nonlinear_stiffness, self.pathway, self.equilibrium_length) + + def test_invalid_constructor_pathway_not_pathway_base(self): + with pytest.raises(TypeError): + _ = DuffingSpring(self.linear_stiffness, self.nonlinear_stiffness, NonSympifyable(), self.equilibrium_length) + + @pytest.mark.parametrize('equilibrium_length', [None, NonSympifyable()]) + def test_invalid_constructor_equilibrium_length_not_sympifyable(self, equilibrium_length): + with pytest.raises(SympifyError): + _ = DuffingSpring(self.linear_stiffness, self.nonlinear_stiffness, self.pathway, equilibrium_length) + + @pytest.mark.parametrize( + 'property_name, fixture_attr_name', + [ + ('linear_stiffness', 'linear_stiffness'), + ('nonlinear_stiffness', 'nonlinear_stiffness'), + ('pathway', 'pathway'), + ('equilibrium_length', 'equilibrium_length') + ] + ) + # Check if certain properties of DuffingSpring object are immutable after initialization + # Ensure that once DuffingSpring is created, its key properties cannot be changed + def test_properties_are_immutable(self, property_name, fixture_attr_name): + spring = DuffingSpring(self.linear_stiffness, self.nonlinear_stiffness, self.pathway, self.equilibrium_length) + with pytest.raises(AttributeError): + setattr(spring, property_name, getattr(self, fixture_attr_name)) + + @pytest.mark.parametrize( + 'equilibrium_length, expected', + [ + (0, 'DuffingSpring(beta, alpha, LinearPathway(pA, pB), equilibrium_length=0)'), + (Symbol('l'), 'DuffingSpring(beta, alpha, LinearPathway(pA, pB), equilibrium_length=l)'), + ] + ) + # Check the __repr__ method of DuffingSpring class + # Check if the actual string representation of DuffingSpring instance matches the expected string for each provided parameter values + def test_repr(self, equilibrium_length, expected): + spring = DuffingSpring(self.linear_stiffness, self.nonlinear_stiffness, self.pathway, equilibrium_length) + assert repr(spring) == expected + + def test_to_loads(self): + self.pB.set_pos(self.pA, self.q*self.N.x) + spring = DuffingSpring(self.linear_stiffness, self.nonlinear_stiffness, self.pathway, self.equilibrium_length) + + # Calculate the displacement from the equilibrium length + displacement = self.q - self.equilibrium_length + + # Make sure this matches the computation in DuffingSpring class + force = -self.linear_stiffness * displacement - self.nonlinear_stiffness * displacement**3 + + # The expected loads on pA and pB due to the spring + expected_loads = [Force(self.pA, force * self.N.x), Force(self.pB, -force * self.N.x)] + + # Compare expected loads to what is returned from DuffingSpring.to_loads() + calculated_loads = spring.to_loads() + for calculated, expected in zip(calculated_loads, expected_loads): + assert calculated.point == expected.point + for dim in self.N: # Assuming self.N is the reference frame + calculated_component = calculated.vector.dot(dim) + expected_component = expected.vector.dot(dim) + # Substitute all symbols with numeric values + substitutions = {self.q: 1, Symbol('l'): 1, Symbol('alpha'): 1, Symbol('beta'): 1} # Add other necessary symbols as needed + diff = (calculated_component - expected_component).subs(substitutions).evalf() + # Check if the absolute value of the difference is below a threshold + assert Abs(diff) < 1e-9, f"The forces do not match. Difference: {diff}" diff --git a/lib/python3.12/site-packages/sympy/physics/mechanics/tests/test_body.py b/lib/python3.12/site-packages/sympy/physics/mechanics/tests/test_body.py new file mode 100644 index 0000000000000000000000000000000000000000..2d59d747400652a0cbb081f4afc5ae4ebaa4db85 --- /dev/null +++ b/lib/python3.12/site-packages/sympy/physics/mechanics/tests/test_body.py @@ -0,0 +1,340 @@ +from sympy import (Symbol, symbols, sin, cos, Matrix, zeros, + simplify) +from sympy.physics.vector import Point, ReferenceFrame, dynamicsymbols, Dyadic +from sympy.physics.mechanics import inertia, Body +from sympy.testing.pytest import raises, warns_deprecated_sympy + + +def test_default(): + with warns_deprecated_sympy(): + body = Body('body') + assert body.name == 'body' + assert body.loads == [] + point = Point('body_masscenter') + point.set_vel(body.frame, 0) + com = body.masscenter + frame = body.frame + assert com.vel(frame) == point.vel(frame) + assert body.mass == Symbol('body_mass') + ixx, iyy, izz = symbols('body_ixx body_iyy body_izz') + ixy, iyz, izx = symbols('body_ixy body_iyz body_izx') + assert body.inertia == (inertia(body.frame, ixx, iyy, izz, ixy, iyz, izx), + body.masscenter) + + +def test_custom_rigid_body(): + # Body with RigidBody. + rigidbody_masscenter = Point('rigidbody_masscenter') + rigidbody_mass = Symbol('rigidbody_mass') + rigidbody_frame = ReferenceFrame('rigidbody_frame') + body_inertia = inertia(rigidbody_frame, 1, 0, 0) + with warns_deprecated_sympy(): + rigid_body = Body('rigidbody_body', rigidbody_masscenter, + rigidbody_mass, rigidbody_frame, body_inertia) + com = rigid_body.masscenter + frame = rigid_body.frame + rigidbody_masscenter.set_vel(rigidbody_frame, 0) + assert com.vel(frame) == rigidbody_masscenter.vel(frame) + assert com.pos_from(com) == rigidbody_masscenter.pos_from(com) + + assert rigid_body.mass == rigidbody_mass + assert rigid_body.inertia == (body_inertia, rigidbody_masscenter) + + assert rigid_body.is_rigidbody + + assert hasattr(rigid_body, 'masscenter') + assert hasattr(rigid_body, 'mass') + assert hasattr(rigid_body, 'frame') + assert hasattr(rigid_body, 'inertia') + + +def test_particle_body(): + # Body with Particle + particle_masscenter = Point('particle_masscenter') + particle_mass = Symbol('particle_mass') + particle_frame = ReferenceFrame('particle_frame') + with warns_deprecated_sympy(): + particle_body = Body('particle_body', particle_masscenter, + particle_mass, particle_frame) + com = particle_body.masscenter + frame = particle_body.frame + particle_masscenter.set_vel(particle_frame, 0) + assert com.vel(frame) == particle_masscenter.vel(frame) + assert com.pos_from(com) == particle_masscenter.pos_from(com) + + assert particle_body.mass == particle_mass + assert not hasattr(particle_body, "_inertia") + assert hasattr(particle_body, 'frame') + assert hasattr(particle_body, 'masscenter') + assert hasattr(particle_body, 'mass') + assert particle_body.inertia == (Dyadic(0), particle_body.masscenter) + assert particle_body.central_inertia == Dyadic(0) + assert not particle_body.is_rigidbody + + particle_body.central_inertia = inertia(particle_frame, 1, 1, 1) + assert particle_body.central_inertia == inertia(particle_frame, 1, 1, 1) + assert particle_body.is_rigidbody + + with warns_deprecated_sympy(): + particle_body = Body('particle_body', mass=particle_mass) + assert not particle_body.is_rigidbody + point = particle_body.masscenter.locatenew('point', particle_body.x) + point_inertia = particle_mass * inertia(particle_body.frame, 0, 1, 1) + particle_body.inertia = (point_inertia, point) + assert particle_body.inertia == (point_inertia, point) + assert particle_body.central_inertia == Dyadic(0) + assert particle_body.is_rigidbody + + +def test_particle_body_add_force(): + # Body with Particle + particle_masscenter = Point('particle_masscenter') + particle_mass = Symbol('particle_mass') + particle_frame = ReferenceFrame('particle_frame') + with warns_deprecated_sympy(): + particle_body = Body('particle_body', particle_masscenter, + particle_mass, particle_frame) + + a = Symbol('a') + force_vector = a * particle_body.frame.x + particle_body.apply_force(force_vector, particle_body.masscenter) + assert len(particle_body.loads) == 1 + point = particle_body.masscenter.locatenew( + particle_body._name + '_point0', 0) + point.set_vel(particle_body.frame, 0) + force_point = particle_body.loads[0][0] + + frame = particle_body.frame + assert force_point.vel(frame) == point.vel(frame) + assert force_point.pos_from(force_point) == point.pos_from(force_point) + + assert particle_body.loads[0][1] == force_vector + + +def test_body_add_force(): + # Body with RigidBody. + rigidbody_masscenter = Point('rigidbody_masscenter') + rigidbody_mass = Symbol('rigidbody_mass') + rigidbody_frame = ReferenceFrame('rigidbody_frame') + body_inertia = inertia(rigidbody_frame, 1, 0, 0) + with warns_deprecated_sympy(): + rigid_body = Body('rigidbody_body', rigidbody_masscenter, + rigidbody_mass, rigidbody_frame, body_inertia) + + l = Symbol('l') + Fa = Symbol('Fa') + point = rigid_body.masscenter.locatenew( + 'rigidbody_body_point0', + l * rigid_body.frame.x) + point.set_vel(rigid_body.frame, 0) + force_vector = Fa * rigid_body.frame.z + # apply_force with point + rigid_body.apply_force(force_vector, point) + assert len(rigid_body.loads) == 1 + force_point = rigid_body.loads[0][0] + frame = rigid_body.frame + assert force_point.vel(frame) == point.vel(frame) + assert force_point.pos_from(force_point) == point.pos_from(force_point) + assert rigid_body.loads[0][1] == force_vector + # apply_force without point + rigid_body.apply_force(force_vector) + assert len(rigid_body.loads) == 2 + assert rigid_body.loads[1][1] == force_vector + # passing something else than point + raises(TypeError, lambda: rigid_body.apply_force(force_vector, 0)) + raises(TypeError, lambda: rigid_body.apply_force(0)) + +def test_body_add_torque(): + with warns_deprecated_sympy(): + body = Body('body') + torque_vector = body.frame.x + body.apply_torque(torque_vector) + + assert len(body.loads) == 1 + assert body.loads[0] == (body.frame, torque_vector) + raises(TypeError, lambda: body.apply_torque(0)) + +def test_body_masscenter_vel(): + with warns_deprecated_sympy(): + A = Body('A') + N = ReferenceFrame('N') + with warns_deprecated_sympy(): + B = Body('B', frame=N) + A.masscenter.set_vel(N, N.z) + assert A.masscenter_vel(B) == N.z + assert A.masscenter_vel(N) == N.z + +def test_body_ang_vel(): + with warns_deprecated_sympy(): + A = Body('A') + N = ReferenceFrame('N') + with warns_deprecated_sympy(): + B = Body('B', frame=N) + A.frame.set_ang_vel(N, N.y) + assert A.ang_vel_in(B) == N.y + assert B.ang_vel_in(A) == -N.y + assert A.ang_vel_in(N) == N.y + +def test_body_dcm(): + with warns_deprecated_sympy(): + A = Body('A') + B = Body('B') + A.frame.orient_axis(B.frame, B.frame.z, 10) + assert A.dcm(B) == Matrix([[cos(10), sin(10), 0], [-sin(10), cos(10), 0], [0, 0, 1]]) + assert A.dcm(B.frame) == Matrix([[cos(10), sin(10), 0], [-sin(10), cos(10), 0], [0, 0, 1]]) + +def test_body_axis(): + N = ReferenceFrame('N') + with warns_deprecated_sympy(): + B = Body('B', frame=N) + assert B.x == N.x + assert B.y == N.y + assert B.z == N.z + +def test_apply_force_multiple_one_point(): + a, b = symbols('a b') + P = Point('P') + with warns_deprecated_sympy(): + B = Body('B') + f1 = a*B.x + f2 = b*B.y + B.apply_force(f1, P) + assert B.loads == [(P, f1)] + B.apply_force(f2, P) + assert B.loads == [(P, f1+f2)] + +def test_apply_force(): + f, g = symbols('f g') + q, x, v1, v2 = dynamicsymbols('q x v1 v2') + P1 = Point('P1') + P2 = Point('P2') + with warns_deprecated_sympy(): + B1 = Body('B1') + B2 = Body('B2') + N = ReferenceFrame('N') + + P1.set_vel(B1.frame, v1*B1.x) + P2.set_vel(B2.frame, v2*B2.x) + force = f*q*N.z # time varying force + + B1.apply_force(force, P1, B2, P2) #applying equal and opposite force on moving points + assert B1.loads == [(P1, force)] + assert B2.loads == [(P2, -force)] + + g1 = B1.mass*g*N.y + g2 = B2.mass*g*N.y + + B1.apply_force(g1) #applying gravity on B1 masscenter + B2.apply_force(g2) #applying gravity on B2 masscenter + + assert B1.loads == [(P1,force), (B1.masscenter, g1)] + assert B2.loads == [(P2, -force), (B2.masscenter, g2)] + + force2 = x*N.x + + B1.apply_force(force2, reaction_body=B2) #Applying time varying force on masscenter + + assert B1.loads == [(P1, force), (B1.masscenter, force2+g1)] + assert B2.loads == [(P2, -force), (B2.masscenter, -force2+g2)] + +def test_apply_torque(): + t = symbols('t') + q = dynamicsymbols('q') + with warns_deprecated_sympy(): + B1 = Body('B1') + B2 = Body('B2') + N = ReferenceFrame('N') + torque = t*q*N.x + + B1.apply_torque(torque, B2) #Applying equal and opposite torque + assert B1.loads == [(B1.frame, torque)] + assert B2.loads == [(B2.frame, -torque)] + + torque2 = t*N.y + B1.apply_torque(torque2) + assert B1.loads == [(B1.frame, torque+torque2)] + +def test_clear_load(): + a = symbols('a') + P = Point('P') + with warns_deprecated_sympy(): + B = Body('B') + force = a*B.z + B.apply_force(force, P) + assert B.loads == [(P, force)] + B.clear_loads() + assert B.loads == [] + +def test_remove_load(): + P1 = Point('P1') + P2 = Point('P2') + with warns_deprecated_sympy(): + B = Body('B') + f1 = B.x + f2 = B.y + B.apply_force(f1, P1) + B.apply_force(f2, P2) + assert B.loads == [(P1, f1), (P2, f2)] + B.remove_load(P2) + assert B.loads == [(P1, f1)] + B.apply_torque(f1.cross(f2)) + assert B.loads == [(P1, f1), (B.frame, f1.cross(f2))] + B.remove_load() + assert B.loads == [(P1, f1)] + +def test_apply_loads_on_multi_degree_freedom_holonomic_system(): + """Example based on: https://pydy.readthedocs.io/en/latest/examples/multidof-holonomic.html""" + with warns_deprecated_sympy(): + W = Body('W') #Wall + B = Body('B') #Block + P = Body('P') #Pendulum + b = Body('b') #bob + q1, q2 = dynamicsymbols('q1 q2') #generalized coordinates + k, c, g, kT = symbols('k c g kT') #constants + F, T = dynamicsymbols('F T') #Specified forces + + #Applying forces + B.apply_force(F*W.x) + W.apply_force(k*q1*W.x, reaction_body=B) #Spring force + W.apply_force(c*q1.diff()*W.x, reaction_body=B) #dampner + P.apply_force(P.mass*g*W.y) + b.apply_force(b.mass*g*W.y) + + #Applying torques + P.apply_torque(kT*q2*W.z, reaction_body=b) + P.apply_torque(T*W.z) + + assert B.loads == [(B.masscenter, (F - k*q1 - c*q1.diff())*W.x)] + assert P.loads == [(P.masscenter, P.mass*g*W.y), (P.frame, (T + kT*q2)*W.z)] + assert b.loads == [(b.masscenter, b.mass*g*W.y), (b.frame, -kT*q2*W.z)] + assert W.loads == [(W.masscenter, (c*q1.diff() + k*q1)*W.x)] + + +def test_parallel_axis(): + N = ReferenceFrame('N') + m, Ix, Iy, Iz, a, b = symbols('m, I_x, I_y, I_z, a, b') + Io = inertia(N, Ix, Iy, Iz) + # Test RigidBody + o = Point('o') + p = o.locatenew('p', a * N.x + b * N.y) + with warns_deprecated_sympy(): + R = Body('R', masscenter=o, frame=N, mass=m, central_inertia=Io) + Ip = R.parallel_axis(p) + Ip_expected = inertia(N, Ix + m * b**2, Iy + m * a**2, + Iz + m * (a**2 + b**2), ixy=-m * a * b) + assert Ip == Ip_expected + # Reference frame from which the parallel axis is viewed should not matter + A = ReferenceFrame('A') + A.orient_axis(N, N.z, 1) + assert simplify( + (R.parallel_axis(p, A) - Ip_expected).to_matrix(A)) == zeros(3, 3) + # Test Particle + o = Point('o') + p = o.locatenew('p', a * N.x + b * N.y) + with warns_deprecated_sympy(): + P = Body('P', masscenter=o, mass=m, frame=N) + Ip = P.parallel_axis(p, N) + Ip_expected = inertia(N, m * b ** 2, m * a ** 2, m * (a ** 2 + b ** 2), + ixy=-m * a * b) + assert not P.is_rigidbody + assert Ip == Ip_expected diff --git a/lib/python3.12/site-packages/sympy/physics/mechanics/tests/test_functions.py b/lib/python3.12/site-packages/sympy/physics/mechanics/tests/test_functions.py new file mode 100644 index 0000000000000000000000000000000000000000..bae6b19b2807dca1632942bd3717e29d214eb269 --- /dev/null +++ b/lib/python3.12/site-packages/sympy/physics/mechanics/tests/test_functions.py @@ -0,0 +1,262 @@ +from sympy import sin, cos, tan, pi, symbols, Matrix, S, Function +from sympy.physics.mechanics import (Particle, Point, ReferenceFrame, + RigidBody) +from sympy.physics.mechanics import (angular_momentum, dynamicsymbols, + kinetic_energy, linear_momentum, + outer, potential_energy, msubs, + find_dynamicsymbols, Lagrangian) + +from sympy.physics.mechanics.functions import ( + center_of_mass, _validate_coordinates, _parse_linear_solver) +from sympy.testing.pytest import raises, warns_deprecated_sympy + + +q1, q2, q3, q4, q5 = symbols('q1 q2 q3 q4 q5') +N = ReferenceFrame('N') +A = N.orientnew('A', 'Axis', [q1, N.z]) +B = A.orientnew('B', 'Axis', [q2, A.x]) +C = B.orientnew('C', 'Axis', [q3, B.y]) + + +def test_linear_momentum(): + N = ReferenceFrame('N') + Ac = Point('Ac') + Ac.set_vel(N, 25 * N.y) + I = outer(N.x, N.x) + A = RigidBody('A', Ac, N, 20, (I, Ac)) + P = Point('P') + Pa = Particle('Pa', P, 1) + Pa.point.set_vel(N, 10 * N.x) + raises(TypeError, lambda: linear_momentum(A, A, Pa)) + raises(TypeError, lambda: linear_momentum(N, N, Pa)) + assert linear_momentum(N, A, Pa) == 10 * N.x + 500 * N.y + + +def test_angular_momentum_and_linear_momentum(): + """A rod with length 2l, centroidal inertia I, and mass M along with a + particle of mass m fixed to the end of the rod rotate with an angular rate + of omega about point O which is fixed to the non-particle end of the rod. + The rod's reference frame is A and the inertial frame is N.""" + m, M, l, I = symbols('m, M, l, I') + omega = dynamicsymbols('omega') + N = ReferenceFrame('N') + a = ReferenceFrame('a') + O = Point('O') + Ac = O.locatenew('Ac', l * N.x) + P = Ac.locatenew('P', l * N.x) + O.set_vel(N, 0 * N.x) + a.set_ang_vel(N, omega * N.z) + Ac.v2pt_theory(O, N, a) + P.v2pt_theory(O, N, a) + Pa = Particle('Pa', P, m) + A = RigidBody('A', Ac, a, M, (I * outer(N.z, N.z), Ac)) + expected = 2 * m * omega * l * N.y + M * l * omega * N.y + assert linear_momentum(N, A, Pa) == expected + raises(TypeError, lambda: angular_momentum(N, N, A, Pa)) + raises(TypeError, lambda: angular_momentum(O, O, A, Pa)) + raises(TypeError, lambda: angular_momentum(O, N, O, Pa)) + expected = (I + M * l**2 + 4 * m * l**2) * omega * N.z + assert angular_momentum(O, N, A, Pa) == expected + + +def test_kinetic_energy(): + m, M, l1 = symbols('m M l1') + omega = dynamicsymbols('omega') + N = ReferenceFrame('N') + O = Point('O') + O.set_vel(N, 0 * N.x) + Ac = O.locatenew('Ac', l1 * N.x) + P = Ac.locatenew('P', l1 * N.x) + a = ReferenceFrame('a') + a.set_ang_vel(N, omega * N.z) + Ac.v2pt_theory(O, N, a) + P.v2pt_theory(O, N, a) + Pa = Particle('Pa', P, m) + I = outer(N.z, N.z) + A = RigidBody('A', Ac, a, M, (I, Ac)) + raises(TypeError, lambda: kinetic_energy(Pa, Pa, A)) + raises(TypeError, lambda: kinetic_energy(N, N, A)) + assert 0 == (kinetic_energy(N, Pa, A) - (M*l1**2*omega**2/2 + + 2*l1**2*m*omega**2 + omega**2/2)).expand() + + +def test_potential_energy(): + m, M, l1, g, h, H = symbols('m M l1 g h H') + omega = dynamicsymbols('omega') + N = ReferenceFrame('N') + O = Point('O') + O.set_vel(N, 0 * N.x) + Ac = O.locatenew('Ac', l1 * N.x) + P = Ac.locatenew('P', l1 * N.x) + a = ReferenceFrame('a') + a.set_ang_vel(N, omega * N.z) + Ac.v2pt_theory(O, N, a) + P.v2pt_theory(O, N, a) + Pa = Particle('Pa', P, m) + I = outer(N.z, N.z) + A = RigidBody('A', Ac, a, M, (I, Ac)) + Pa.potential_energy = m * g * h + A.potential_energy = M * g * H + assert potential_energy(A, Pa) == m * g * h + M * g * H + + +def test_Lagrangian(): + M, m, g, h = symbols('M m g h') + N = ReferenceFrame('N') + O = Point('O') + O.set_vel(N, 0 * N.x) + P = O.locatenew('P', 1 * N.x) + P.set_vel(N, 10 * N.x) + Pa = Particle('Pa', P, 1) + Ac = O.locatenew('Ac', 2 * N.y) + Ac.set_vel(N, 5 * N.y) + a = ReferenceFrame('a') + a.set_ang_vel(N, 10 * N.z) + I = outer(N.z, N.z) + A = RigidBody('A', Ac, a, 20, (I, Ac)) + Pa.potential_energy = m * g * h + A.potential_energy = M * g * h + raises(TypeError, lambda: Lagrangian(A, A, Pa)) + raises(TypeError, lambda: Lagrangian(N, N, Pa)) + + +def test_msubs(): + a, b = symbols('a, b') + x, y, z = dynamicsymbols('x, y, z') + # Test simple substitution + expr = Matrix([[a*x + b, x*y.diff() + y], + [x.diff().diff(), z + sin(z.diff())]]) + sol = Matrix([[a + b, y], + [x.diff().diff(), 1]]) + sd = {x: 1, z: 1, z.diff(): 0, y.diff(): 0} + assert msubs(expr, sd) == sol + # Test smart substitution + expr = cos(x + y)*tan(x + y) + b*x.diff() + sd = {x: 0, y: pi/2, x.diff(): 1} + assert msubs(expr, sd, smart=True) == b + 1 + N = ReferenceFrame('N') + v = x*N.x + y*N.y + d = x*(N.x|N.x) + y*(N.y|N.y) + v_sol = 1*N.y + d_sol = 1*(N.y|N.y) + sd = {x: 0, y: 1} + assert msubs(v, sd) == v_sol + assert msubs(d, sd) == d_sol + + +def test_find_dynamicsymbols(): + a, b = symbols('a, b') + x, y, z = dynamicsymbols('x, y, z') + expr = Matrix([[a*x + b, x*y.diff() + y], + [x.diff().diff(), z + sin(z.diff())]]) + # Test finding all dynamicsymbols + sol = {x, y.diff(), y, x.diff().diff(), z, z.diff()} + assert find_dynamicsymbols(expr) == sol + # Test finding all but those in sym_list + exclude_list = [x, y, z] + sol = {y.diff(), x.diff().diff(), z.diff()} + assert find_dynamicsymbols(expr, exclude=exclude_list) == sol + # Test finding all dynamicsymbols in a vector with a given reference frame + d, e, f = dynamicsymbols('d, e, f') + A = ReferenceFrame('A') + v = d * A.x + e * A.y + f * A.z + sol = {d, e, f} + assert find_dynamicsymbols(v, reference_frame=A) == sol + # Test if a ValueError is raised on supplying only a vector as input + raises(ValueError, lambda: find_dynamicsymbols(v)) + + +# This function tests the center_of_mass() function +# that was added in PR #14758 to compute the center of +# mass of a system of bodies. +def test_center_of_mass(): + a = ReferenceFrame('a') + m = symbols('m', real=True) + p1 = Particle('p1', Point('p1_pt'), S.One) + p2 = Particle('p2', Point('p2_pt'), S(2)) + p3 = Particle('p3', Point('p3_pt'), S(3)) + p4 = Particle('p4', Point('p4_pt'), m) + b_f = ReferenceFrame('b_f') + b_cm = Point('b_cm') + mb = symbols('mb') + b = RigidBody('b', b_cm, b_f, mb, (outer(b_f.x, b_f.x), b_cm)) + p2.point.set_pos(p1.point, a.x) + p3.point.set_pos(p1.point, a.x + a.y) + p4.point.set_pos(p1.point, a.y) + b.masscenter.set_pos(p1.point, a.y + a.z) + point_o=Point('o') + point_o.set_pos(p1.point, center_of_mass(p1.point, p1, p2, p3, p4, b)) + expr = 5/(m + mb + 6)*a.x + (m + mb + 3)/(m + mb + 6)*a.y + mb/(m + mb + 6)*a.z + assert point_o.pos_from(p1.point)-expr == 0 + + +def test_validate_coordinates(): + q1, q2, q3, u1, u2, u3, ua1, ua2, ua3 = dynamicsymbols('q1:4 u1:4 ua1:4') + s1, s2, s3 = symbols('s1:4') + # Test normal + _validate_coordinates([q1, q2, q3], [u1, u2, u3], + u_auxiliary=[ua1, ua2, ua3]) + # Test not equal number of coordinates and speeds + _validate_coordinates([q1, q2]) + _validate_coordinates([q1, q2], [u1]) + _validate_coordinates(speeds=[u1, u2]) + # Test duplicate + _validate_coordinates([q1, q2, q2], [u1, u2, u3], check_duplicates=False) + raises(ValueError, lambda: _validate_coordinates( + [q1, q2, q2], [u1, u2, u3])) + _validate_coordinates([q1, q2, q3], [u1, u2, u2], check_duplicates=False) + raises(ValueError, lambda: _validate_coordinates( + [q1, q2, q3], [u1, u2, u2], check_duplicates=True)) + raises(ValueError, lambda: _validate_coordinates( + [q1, q2, q3], [q1, u2, u3], check_duplicates=True)) + _validate_coordinates([q1, q2, q3], [u1, u2, u3], check_duplicates=False, + u_auxiliary=[u1, ua2, ua2]) + raises(ValueError, lambda: _validate_coordinates( + [q1, q2, q3], [u1, u2, u3], u_auxiliary=[u1, ua2, ua3])) + raises(ValueError, lambda: _validate_coordinates( + [q1, q2, q3], [u1, u2, u3], u_auxiliary=[q1, ua2, ua3])) + raises(ValueError, lambda: _validate_coordinates( + [q1, q2, q3], [u1, u2, u3], u_auxiliary=[ua1, ua2, ua2])) + # Test is_dynamicsymbols + _validate_coordinates([q1 + q2, q3], is_dynamicsymbols=False) + raises(ValueError, lambda: _validate_coordinates([q1 + q2, q3])) + _validate_coordinates([s1, q1, q2], [0, u1, u2], is_dynamicsymbols=False) + raises(ValueError, lambda: _validate_coordinates( + [s1, q1, q2], [0, u1, u2], is_dynamicsymbols=True)) + _validate_coordinates([s1 + s2 + s3, q1], [0, u1], is_dynamicsymbols=False) + raises(ValueError, lambda: _validate_coordinates( + [s1 + s2 + s3, q1], [0, u1], is_dynamicsymbols=True)) + _validate_coordinates(u_auxiliary=[s1, ua1], is_dynamicsymbols=False) + raises(ValueError, lambda: _validate_coordinates(u_auxiliary=[s1, ua1])) + # Test normal function + t = dynamicsymbols._t + a = symbols('a') + f1, f2 = symbols('f1:3', cls=Function) + _validate_coordinates([f1(a), f2(a)], is_dynamicsymbols=False) + raises(ValueError, lambda: _validate_coordinates([f1(a), f2(a)])) + raises(ValueError, lambda: _validate_coordinates(speeds=[f1(a), f2(a)])) + dynamicsymbols._t = a + _validate_coordinates([f1(a), f2(a)]) + raises(ValueError, lambda: _validate_coordinates([f1(t), f2(t)])) + dynamicsymbols._t = t + + +def test_parse_linear_solver(): + A, b = Matrix(3, 3, symbols('a:9')), Matrix(3, 2, symbols('b:6')) + assert _parse_linear_solver(Matrix.LUsolve) == Matrix.LUsolve # Test callable + assert _parse_linear_solver('LU')(A, b) == Matrix.LUsolve(A, b) + + +def test_deprecated_moved_functions(): + from sympy.physics.mechanics.functions import ( + inertia, inertia_of_point_mass, gravity) + N = ReferenceFrame('N') + with warns_deprecated_sympy(): + assert inertia(N, 0, 1, 0, 1) == (N.x | N.y) + (N.y | N.x) + (N.y | N.y) + with warns_deprecated_sympy(): + assert inertia_of_point_mass(1, N.x + N.y, N) == ( + (N.x | N.x) + (N.y | N.y) + 2 * (N.z | N.z) - + (N.x | N.y) - (N.y | N.x)) + p = Particle('P') + with warns_deprecated_sympy(): + assert gravity(-2 * N.z, p) == [(p.masscenter, -2 * p.mass * N.z)] diff --git a/lib/python3.12/site-packages/sympy/physics/mechanics/tests/test_inertia.py b/lib/python3.12/site-packages/sympy/physics/mechanics/tests/test_inertia.py new file mode 100644 index 0000000000000000000000000000000000000000..8d29e5f31868e539c4b50575af5180e5eb96f2cd --- /dev/null +++ b/lib/python3.12/site-packages/sympy/physics/mechanics/tests/test_inertia.py @@ -0,0 +1,71 @@ +from sympy import symbols +from sympy.testing.pytest import raises +from sympy.physics.mechanics import (inertia, inertia_of_point_mass, + Inertia, ReferenceFrame, Point) + + +def test_inertia_dyadic(): + N = ReferenceFrame('N') + ixx, iyy, izz = symbols('ixx iyy izz') + ixy, iyz, izx = symbols('ixy iyz izx') + assert inertia(N, ixx, iyy, izz) == (ixx * (N.x | N.x) + iyy * + (N.y | N.y) + izz * (N.z | N.z)) + assert inertia(N, 0, 0, 0) == 0 * (N.x | N.x) + raises(TypeError, lambda: inertia(0, 0, 0, 0)) + assert inertia(N, ixx, iyy, izz, ixy, iyz, izx) == (ixx * (N.x | N.x) + + ixy * (N.x | N.y) + izx * (N.x | N.z) + ixy * (N.y | N.x) + iyy * + (N.y | N.y) + iyz * (N.y | N.z) + izx * (N.z | N.x) + iyz * (N.z | + N.y) + izz * (N.z | N.z)) + + +def test_inertia_of_point_mass(): + r, s, t, m = symbols('r s t m') + N = ReferenceFrame('N') + + px = r * N.x + I = inertia_of_point_mass(m, px, N) + assert I == m * r**2 * (N.y | N.y) + m * r**2 * (N.z | N.z) + + py = s * N.y + I = inertia_of_point_mass(m, py, N) + assert I == m * s**2 * (N.x | N.x) + m * s**2 * (N.z | N.z) + + pz = t * N.z + I = inertia_of_point_mass(m, pz, N) + assert I == m * t**2 * (N.x | N.x) + m * t**2 * (N.y | N.y) + + p = px + py + pz + I = inertia_of_point_mass(m, p, N) + assert I == (m * (s**2 + t**2) * (N.x | N.x) - + m * r * s * (N.x | N.y) - + m * r * t * (N.x | N.z) - + m * r * s * (N.y | N.x) + + m * (r**2 + t**2) * (N.y | N.y) - + m * s * t * (N.y | N.z) - + m * r * t * (N.z | N.x) - + m * s * t * (N.z | N.y) + + m * (r**2 + s**2) * (N.z | N.z)) + + +def test_inertia_object(): + N = ReferenceFrame('N') + O = Point('O') + ixx, iyy, izz = symbols('ixx iyy izz') + I_dyadic = ixx * (N.x | N.x) + iyy * (N.y | N.y) + izz * (N.z | N.z) + I = Inertia(inertia(N, ixx, iyy, izz), O) + assert isinstance(I, tuple) + assert I.__repr__() == ('Inertia(dyadic=ixx*(N.x|N.x) + iyy*(N.y|N.y) + ' + 'izz*(N.z|N.z), point=O)') + assert I.dyadic == I_dyadic + assert I.point == O + assert I[0] == I_dyadic + assert I[1] == O + assert I == (I_dyadic, O) # Test tuple equal + raises(TypeError, lambda: I != (O, I_dyadic)) # Incorrect tuple order + assert I == Inertia(O, I_dyadic) # Parse changed argument order + assert I == Inertia.from_inertia_scalars(O, N, ixx, iyy, izz) + # Test invalid tuple operations + raises(TypeError, lambda: I + (1, 2)) + raises(TypeError, lambda: (1, 2) + I) + raises(TypeError, lambda: I * 2) + raises(TypeError, lambda: 2 * I) diff --git a/lib/python3.12/site-packages/sympy/physics/mechanics/tests/test_joint.py b/lib/python3.12/site-packages/sympy/physics/mechanics/tests/test_joint.py new file mode 100644 index 0000000000000000000000000000000000000000..271801b5b7290a4479ee61e1414741e4c4d6f966 --- /dev/null +++ b/lib/python3.12/site-packages/sympy/physics/mechanics/tests/test_joint.py @@ -0,0 +1,1240 @@ +from sympy.core.function import expand_mul +from sympy.core.numbers import pi +from sympy.core.singleton import S +from sympy.functions.elementary.miscellaneous import sqrt +from sympy.functions.elementary.trigonometric import (cos, sin) +from sympy import Matrix, simplify, eye, zeros +from sympy.core.symbol import symbols +from sympy.physics.mechanics import ( + dynamicsymbols, RigidBody, Particle, JointsMethod, PinJoint, PrismaticJoint, + CylindricalJoint, PlanarJoint, SphericalJoint, WeldJoint, Body) +from sympy.physics.mechanics.joint import Joint +from sympy.physics.vector import Vector, ReferenceFrame, Point +from sympy.testing.pytest import raises, warns_deprecated_sympy + + +t = dynamicsymbols._t # type: ignore + + +def _generate_body(interframe=False): + N = ReferenceFrame('N') + A = ReferenceFrame('A') + P = RigidBody('P', frame=N) + C = RigidBody('C', frame=A) + if interframe: + Pint, Cint = ReferenceFrame('P_int'), ReferenceFrame('C_int') + Pint.orient_axis(N, N.x, pi) + Cint.orient_axis(A, A.y, -pi / 2) + return N, A, P, C, Pint, Cint + return N, A, P, C + + +def test_Joint(): + parent = RigidBody('parent') + child = RigidBody('child') + raises(TypeError, lambda: Joint('J', parent, child)) + + +def test_coordinate_generation(): + q, u, qj, uj = dynamicsymbols('q u q_J u_J') + q0j, q1j, q2j, q3j, u0j, u1j, u2j, u3j = dynamicsymbols('q0:4_J u0:4_J') + q0, q1, q2, q3, u0, u1, u2, u3 = dynamicsymbols('q0:4 u0:4') + _, _, P, C = _generate_body() + # Using PinJoint to access Joint's coordinate generation method + J = PinJoint('J', P, C) + # Test single given + assert J._fill_coordinate_list(q, 1) == Matrix([q]) + assert J._fill_coordinate_list([u], 1) == Matrix([u]) + assert J._fill_coordinate_list([u], 1, offset=2) == Matrix([u]) + # Test None + assert J._fill_coordinate_list(None, 1) == Matrix([qj]) + assert J._fill_coordinate_list([None], 1) == Matrix([qj]) + assert J._fill_coordinate_list([q0, None, None], 3) == Matrix( + [q0, q1j, q2j]) + # Test autofill + assert J._fill_coordinate_list(None, 3) == Matrix([q0j, q1j, q2j]) + assert J._fill_coordinate_list([], 3) == Matrix([q0j, q1j, q2j]) + # Test offset + assert J._fill_coordinate_list([], 3, offset=1) == Matrix([q1j, q2j, q3j]) + assert J._fill_coordinate_list([q1, None, q3], 3, offset=1) == Matrix( + [q1, q2j, q3]) + assert J._fill_coordinate_list(None, 2, offset=2) == Matrix([q2j, q3j]) + # Test label + assert J._fill_coordinate_list(None, 1, 'u') == Matrix([uj]) + assert J._fill_coordinate_list([], 3, 'u') == Matrix([u0j, u1j, u2j]) + # Test single numbering + assert J._fill_coordinate_list(None, 1, number_single=True) == Matrix([q0j]) + assert J._fill_coordinate_list([], 1, 'u', 2, True) == Matrix([u2j]) + assert J._fill_coordinate_list([], 3, 'q') == Matrix([q0j, q1j, q2j]) + # Test invalid number of coordinates supplied + raises(ValueError, lambda: J._fill_coordinate_list([q0, q1], 1)) + raises(ValueError, lambda: J._fill_coordinate_list([u0, u1, None], 2, 'u')) + raises(ValueError, lambda: J._fill_coordinate_list([q0, q1], 3)) + # Test incorrect coordinate type + raises(TypeError, lambda: J._fill_coordinate_list([q0, symbols('q1')], 2)) + raises(TypeError, lambda: J._fill_coordinate_list([q0 + q1, q1], 2)) + # Test if derivative as generalized speed is allowed + _, _, P, C = _generate_body() + PinJoint('J', P, C, q1, q1.diff(t)) + # Test duplicate coordinates + _, _, P, C = _generate_body() + raises(ValueError, lambda: SphericalJoint('J', P, C, [q1j, None, None])) + raises(ValueError, lambda: SphericalJoint('J', P, C, speeds=[u0, u0, u1])) + + +def test_pin_joint(): + P = RigidBody('P') + C = RigidBody('C') + l, m = symbols('l m') + q, u = dynamicsymbols('q_J, u_J') + Pj = PinJoint('J', P, C) + assert Pj.name == 'J' + assert Pj.parent == P + assert Pj.child == C + assert Pj.coordinates == Matrix([q]) + assert Pj.speeds == Matrix([u]) + assert Pj.kdes == Matrix([u - q.diff(t)]) + assert Pj.joint_axis == P.frame.x + assert Pj.child_point.pos_from(C.masscenter) == Vector(0) + assert Pj.parent_point.pos_from(P.masscenter) == Vector(0) + assert Pj.parent_point.pos_from(Pj._child_point) == Vector(0) + assert C.masscenter.pos_from(P.masscenter) == Vector(0) + assert Pj.parent_interframe == P.frame + assert Pj.child_interframe == C.frame + assert Pj.__str__() == 'PinJoint: J parent: P child: C' + + P1 = RigidBody('P1') + C1 = RigidBody('C1') + Pint = ReferenceFrame('P_int') + Pint.orient_axis(P1.frame, P1.y, pi / 2) + J1 = PinJoint('J1', P1, C1, parent_point=l*P1.frame.x, + child_point=m*C1.frame.y, joint_axis=P1.frame.z, + parent_interframe=Pint) + assert J1._joint_axis == P1.frame.z + assert J1._child_point.pos_from(C1.masscenter) == m * C1.frame.y + assert J1._parent_point.pos_from(P1.masscenter) == l * P1.frame.x + assert J1._parent_point.pos_from(J1._child_point) == Vector(0) + assert (P1.masscenter.pos_from(C1.masscenter) == + -l*P1.frame.x + m*C1.frame.y) + assert J1.parent_interframe == Pint + assert J1.child_interframe == C1.frame + + q, u = dynamicsymbols('q, u') + N, A, P, C, Pint, Cint = _generate_body(True) + parent_point = P.masscenter.locatenew('parent_point', N.x + N.y) + child_point = C.masscenter.locatenew('child_point', C.y + C.z) + J = PinJoint('J', P, C, q, u, parent_point=parent_point, + child_point=child_point, parent_interframe=Pint, + child_interframe=Cint, joint_axis=N.z) + assert J.joint_axis == N.z + assert J.parent_point.vel(N) == 0 + assert J.parent_point == parent_point + assert J.child_point == child_point + assert J.child_point.pos_from(P.masscenter) == N.x + N.y + assert J.parent_point.pos_from(C.masscenter) == C.y + C.z + assert C.masscenter.pos_from(P.masscenter) == N.x + N.y - C.y - C.z + assert C.masscenter.vel(N).express(N) == (u * sin(q) - u * cos(q)) * N.x + ( + -u * sin(q) - u * cos(q)) * N.y + assert J.parent_interframe == Pint + assert J.child_interframe == Cint + + +def test_particle_compatibility(): + m, l = symbols('m l') + C_frame = ReferenceFrame('C') + P = Particle('P') + C = Particle('C', mass=m) + q, u = dynamicsymbols('q, u') + J = PinJoint('J', P, C, q, u, child_interframe=C_frame, + child_point=l * C_frame.y) + assert J.child_interframe == C_frame + assert J.parent_interframe.name == 'J_P_frame' + assert C.masscenter.pos_from(P.masscenter) == -l * C_frame.y + assert C_frame.dcm(J.parent_interframe) == Matrix([[1, 0, 0], + [0, cos(q), sin(q)], + [0, -sin(q), cos(q)]]) + assert C.masscenter.vel(J.parent_interframe) == -l * u * C_frame.z + # Test with specified joint axis + P_frame = ReferenceFrame('P') + C_frame = ReferenceFrame('C') + P = Particle('P') + C = Particle('C', mass=m) + q, u = dynamicsymbols('q, u') + J = PinJoint('J', P, C, q, u, parent_interframe=P_frame, + child_interframe=C_frame, child_point=l * C_frame.y, + joint_axis=P_frame.z) + assert J.joint_axis == J.parent_interframe.z + assert C_frame.dcm(J.parent_interframe) == Matrix([[cos(q), sin(q), 0], + [-sin(q), cos(q), 0], + [0, 0, 1]]) + assert P.masscenter.vel(J.parent_interframe) == 0 + assert C.masscenter.vel(J.parent_interframe) == l * u * C_frame.x + q1, q2, q3, u1, u2, u3 = dynamicsymbols('q1:4 u1:4') + qdot_to_u = {qi.diff(t): ui for qi, ui in ((q1, u1), (q2, u2), (q3, u3))} + # Test compatibility for prismatic joint + P, C = Particle('P'), Particle('C') + J = PrismaticJoint('J', P, C, q, u) + assert J.parent_interframe.dcm(J.child_interframe) == eye(3) + assert C.masscenter.pos_from(P.masscenter) == q * J.parent_interframe.x + assert P.masscenter.vel(J.parent_interframe) == 0 + assert C.masscenter.vel(J.parent_interframe) == u * J.parent_interframe.x + # Test compatibility for cylindrical joint + P, C = Particle('P'), Particle('C') + P_frame = ReferenceFrame('P_frame') + J = CylindricalJoint('J', P, C, q1, q2, u1, u2, parent_interframe=P_frame, + parent_point=l * P_frame.x, joint_axis=P_frame.y) + assert J.parent_interframe.dcm(J.child_interframe) == Matrix([ + [cos(q1), 0, sin(q1)], [0, 1, 0], [-sin(q1), 0, cos(q1)]]) + assert C.masscenter.pos_from(P.masscenter) == l * P_frame.x + q2 * P_frame.y + assert C.masscenter.vel(J.parent_interframe) == u2 * P_frame.y + assert P.masscenter.vel(J.child_interframe).xreplace(qdot_to_u) == ( + -u2 * P_frame.y - l * u1 * P_frame.z) + # Test compatibility for planar joint + P, C = Particle('P'), Particle('C') + C_frame = ReferenceFrame('C_frame') + J = PlanarJoint('J', P, C, q1, [q2, q3], u1, [u2, u3], + child_interframe=C_frame, child_point=l * C_frame.z) + P_frame = J.parent_interframe + assert J.parent_interframe.dcm(J.child_interframe) == Matrix([ + [1, 0, 0], [0, cos(q1), -sin(q1)], [0, sin(q1), cos(q1)]]) + assert C.masscenter.pos_from(P.masscenter) == ( + -l * C_frame.z + q2 * P_frame.y + q3 * P_frame.z) + assert C.masscenter.vel(J.parent_interframe) == ( + l * u1 * C_frame.y + u2 * P_frame.y + u3 * P_frame.z) + # Test compatibility for weld joint + P, C = Particle('P'), Particle('C') + C_frame, P_frame = ReferenceFrame('C_frame'), ReferenceFrame('P_frame') + J = WeldJoint('J', P, C, parent_interframe=P_frame, + child_interframe=C_frame, parent_point=l * P_frame.x, + child_point=l * C_frame.y) + assert P_frame.dcm(C_frame) == eye(3) + assert C.masscenter.pos_from(P.masscenter) == l * P_frame.x - l * C_frame.y + assert C.masscenter.vel(J.parent_interframe) == 0 + + +def test_body_compatibility(): + m, l = symbols('m l') + C_frame = ReferenceFrame('C') + with warns_deprecated_sympy(): + P = Body('P') + C = Body('C', mass=m, frame=C_frame) + q, u = dynamicsymbols('q, u') + PinJoint('J', P, C, q, u, child_point=l * C_frame.y) + assert C.frame == C_frame + assert P.frame.name == 'P_frame' + assert C.masscenter.pos_from(P.masscenter) == -l * C.y + assert C.frame.dcm(P.frame) == Matrix([[1, 0, 0], + [0, cos(q), sin(q)], + [0, -sin(q), cos(q)]]) + assert C.masscenter.vel(P.frame) == -l * u * C.z + + +def test_pin_joint_double_pendulum(): + q1, q2 = dynamicsymbols('q1 q2') + u1, u2 = dynamicsymbols('u1 u2') + m, l = symbols('m l') + N = ReferenceFrame('N') + A = ReferenceFrame('A') + B = ReferenceFrame('B') + C = RigidBody('C', frame=N) # ceiling + PartP = RigidBody('P', frame=A, mass=m) + PartR = RigidBody('R', frame=B, mass=m) + + J1 = PinJoint('J1', C, PartP, speeds=u1, coordinates=q1, + child_point=-l*A.x, joint_axis=C.frame.z) + J2 = PinJoint('J2', PartP, PartR, speeds=u2, coordinates=q2, + child_point=-l*B.x, joint_axis=PartP.frame.z) + + # Check orientation + assert N.dcm(A) == Matrix([[cos(q1), -sin(q1), 0], + [sin(q1), cos(q1), 0], [0, 0, 1]]) + assert A.dcm(B) == Matrix([[cos(q2), -sin(q2), 0], + [sin(q2), cos(q2), 0], [0, 0, 1]]) + assert simplify(N.dcm(B)) == Matrix([[cos(q1 + q2), -sin(q1 + q2), 0], + [sin(q1 + q2), cos(q1 + q2), 0], + [0, 0, 1]]) + + # Check Angular Velocity + assert A.ang_vel_in(N) == u1 * N.z + assert B.ang_vel_in(A) == u2 * A.z + assert B.ang_vel_in(N) == u1 * N.z + u2 * A.z + + # Check kde + assert J1.kdes == Matrix([u1 - q1.diff(t)]) + assert J2.kdes == Matrix([u2 - q2.diff(t)]) + + # Check Linear Velocity + assert PartP.masscenter.vel(N) == l*u1*A.y + assert PartR.masscenter.vel(A) == l*u2*B.y + assert PartR.masscenter.vel(N) == l*u1*A.y + l*(u1 + u2)*B.y + + +def test_pin_joint_chaos_pendulum(): + mA, mB, lA, lB, h = symbols('mA, mB, lA, lB, h') + theta, phi, omega, alpha = dynamicsymbols('theta phi omega alpha') + N = ReferenceFrame('N') + A = ReferenceFrame('A') + B = ReferenceFrame('B') + lA = (lB - h / 2) / 2 + lC = (lB/2 + h/4) + rod = RigidBody('rod', frame=A, mass=mA) + plate = RigidBody('plate', mass=mB, frame=B) + C = RigidBody('C', frame=N) + J1 = PinJoint('J1', C, rod, coordinates=theta, speeds=omega, + child_point=lA*A.z, joint_axis=N.y) + J2 = PinJoint('J2', rod, plate, coordinates=phi, speeds=alpha, + parent_point=lC*A.z, joint_axis=A.z) + + # Check orientation + assert A.dcm(N) == Matrix([[cos(theta), 0, -sin(theta)], + [0, 1, 0], + [sin(theta), 0, cos(theta)]]) + assert A.dcm(B) == Matrix([[cos(phi), -sin(phi), 0], + [sin(phi), cos(phi), 0], + [0, 0, 1]]) + assert B.dcm(N) == Matrix([ + [cos(phi)*cos(theta), sin(phi), -sin(theta)*cos(phi)], + [-sin(phi)*cos(theta), cos(phi), sin(phi)*sin(theta)], + [sin(theta), 0, cos(theta)]]) + + # Check Angular Velocity + assert A.ang_vel_in(N) == omega*N.y + assert A.ang_vel_in(B) == -alpha*A.z + assert N.ang_vel_in(B) == -omega*N.y - alpha*A.z + + # Check kde + assert J1.kdes == Matrix([omega - theta.diff(t)]) + assert J2.kdes == Matrix([alpha - phi.diff(t)]) + + # Check pos of masscenters + assert C.masscenter.pos_from(rod.masscenter) == lA*A.z + assert rod.masscenter.pos_from(plate.masscenter) == - lC * A.z + + # Check Linear Velocities + assert rod.masscenter.vel(N) == (h/4 - lB/2)*omega*A.x + assert plate.masscenter.vel(N) == ((h/4 - lB/2)*omega + + (h/4 + lB/2)*omega)*A.x + + +def test_pin_joint_interframe(): + q, u = dynamicsymbols('q, u') + # Check not connected + N, A, P, C = _generate_body() + Pint, Cint = ReferenceFrame('Pint'), ReferenceFrame('Cint') + raises(ValueError, lambda: PinJoint('J', P, C, parent_interframe=Pint)) + raises(ValueError, lambda: PinJoint('J', P, C, child_interframe=Cint)) + # Check not fixed interframe + Pint.orient_axis(N, N.z, q) + Cint.orient_axis(A, A.z, q) + raises(ValueError, lambda: PinJoint('J', P, C, parent_interframe=Pint)) + raises(ValueError, lambda: PinJoint('J', P, C, child_interframe=Cint)) + # Check only parent_interframe + N, A, P, C = _generate_body() + Pint = ReferenceFrame('Pint') + Pint.orient_body_fixed(N, (pi / 4, pi, pi / 3), 'xyz') + PinJoint('J', P, C, q, u, parent_point=N.x, child_point=-C.y, + parent_interframe=Pint, joint_axis=Pint.x) + assert simplify(N.dcm(A)) - Matrix([ + [-1 / 2, sqrt(3) * cos(q) / 2, -sqrt(3) * sin(q) / 2], + [sqrt(6) / 4, sqrt(2) * (2 * sin(q) + cos(q)) / 4, + sqrt(2) * (-sin(q) + 2 * cos(q)) / 4], + [sqrt(6) / 4, sqrt(2) * (-2 * sin(q) + cos(q)) / 4, + -sqrt(2) * (sin(q) + 2 * cos(q)) / 4]]) == zeros(3) + assert A.ang_vel_in(N) == u * Pint.x + assert C.masscenter.pos_from(P.masscenter) == N.x + A.y + assert C.masscenter.vel(N) == u * A.z + assert P.masscenter.vel(Pint) == Vector(0) + assert C.masscenter.vel(Pint) == u * A.z + # Check only child_interframe + N, A, P, C = _generate_body() + Cint = ReferenceFrame('Cint') + Cint.orient_body_fixed(A, (2 * pi / 3, -pi, pi / 2), 'xyz') + PinJoint('J', P, C, q, u, parent_point=-N.z, child_point=C.x, + child_interframe=Cint, joint_axis=P.x + P.z) + assert simplify(N.dcm(A)) == Matrix([ + [-sqrt(2) * sin(q) / 2, + -sqrt(3) * (cos(q) - 1) / 4 - cos(q) / 4 - S(1) / 4, + sqrt(3) * (cos(q) + 1) / 4 - cos(q) / 4 + S(1) / 4], + [cos(q), (sqrt(2) + sqrt(6)) * -sin(q) / 4, + (-sqrt(2) + sqrt(6)) * sin(q) / 4], + [sqrt(2) * sin(q) / 2, + sqrt(3) * (cos(q) + 1) / 4 + cos(q) / 4 - S(1) / 4, + sqrt(3) * (1 - cos(q)) / 4 + cos(q) / 4 + S(1) / 4]]) + assert A.ang_vel_in(N) == sqrt(2) * u / 2 * N.x + sqrt(2) * u / 2 * N.z + assert C.masscenter.pos_from(P.masscenter) == - N.z - A.x + assert C.masscenter.vel(N).simplify() == ( + -sqrt(6) - sqrt(2)) * u / 4 * A.y + ( + -sqrt(2) + sqrt(6)) * u / 4 * A.z + assert C.masscenter.vel(Cint) == Vector(0) + # Check combination + N, A, P, C = _generate_body() + Pint, Cint = ReferenceFrame('Pint'), ReferenceFrame('Cint') + Pint.orient_body_fixed(N, (-pi / 2, pi, pi / 2), 'xyz') + Cint.orient_body_fixed(A, (2 * pi / 3, -pi, pi / 2), 'xyz') + PinJoint('J', P, C, q, u, parent_point=N.x - N.y, child_point=-C.z, + parent_interframe=Pint, child_interframe=Cint, + joint_axis=Pint.x + Pint.z) + assert simplify(N.dcm(A)) == Matrix([ + [cos(q), (sqrt(2) + sqrt(6)) * -sin(q) / 4, + (-sqrt(2) + sqrt(6)) * sin(q) / 4], + [-sqrt(2) * sin(q) / 2, + -sqrt(3) * (cos(q) + 1) / 4 - cos(q) / 4 + S(1) / 4, + sqrt(3) * (cos(q) - 1) / 4 - cos(q) / 4 - S(1) / 4], + [sqrt(2) * sin(q) / 2, + sqrt(3) * (cos(q) - 1) / 4 + cos(q) / 4 + S(1) / 4, + -sqrt(3) * (cos(q) + 1) / 4 + cos(q) / 4 - S(1) / 4]]) + assert A.ang_vel_in(N) == sqrt(2) * u / 2 * Pint.x + sqrt( + 2) * u / 2 * Pint.z + assert C.masscenter.pos_from(P.masscenter) == N.x - N.y + A.z + N_v_C = (-sqrt(2) + sqrt(6)) * u / 4 * A.x + assert C.masscenter.vel(N).simplify() == N_v_C + assert C.masscenter.vel(Pint).simplify() == N_v_C + assert C.masscenter.vel(Cint) == Vector(0) + + +def test_pin_joint_joint_axis(): + q, u = dynamicsymbols('q, u') + # Check parent as reference + N, A, P, C, Pint, Cint = _generate_body(True) + pin = PinJoint('J', P, C, q, u, parent_interframe=Pint, + child_interframe=Cint, joint_axis=P.y) + assert pin.joint_axis == P.y + assert N.dcm(A) == Matrix([[sin(q), 0, cos(q)], [0, -1, 0], + [cos(q), 0, -sin(q)]]) + # Check parent_interframe as reference + N, A, P, C, Pint, Cint = _generate_body(True) + pin = PinJoint('J', P, C, q, u, parent_interframe=Pint, + child_interframe=Cint, joint_axis=Pint.y) + assert pin.joint_axis == Pint.y + assert N.dcm(A) == Matrix([[-sin(q), 0, cos(q)], [0, -1, 0], + [cos(q), 0, sin(q)]]) + # Check combination of joint_axis with interframes supplied as vectors (2x) + N, A, P, C = _generate_body() + pin = PinJoint('J', P, C, q, u, parent_interframe=N.z, + child_interframe=-C.z, joint_axis=N.z) + assert pin.joint_axis == N.z + assert N.dcm(A) == Matrix([[-cos(q), -sin(q), 0], [-sin(q), cos(q), 0], + [0, 0, -1]]) + N, A, P, C = _generate_body() + pin = PinJoint('J', P, C, q, u, parent_interframe=N.z, + child_interframe=-C.z, joint_axis=N.x) + assert pin.joint_axis == N.x + assert N.dcm(A) == Matrix([[-1, 0, 0], [0, cos(q), sin(q)], + [0, sin(q), -cos(q)]]) + # Check time varying axis + N, A, P, C, Pint, Cint = _generate_body(True) + raises(ValueError, lambda: PinJoint('J', P, C, + joint_axis=cos(q) * N.x + sin(q) * N.y)) + # Check joint_axis provided in child frame + raises(ValueError, lambda: PinJoint('J', P, C, joint_axis=C.x)) + # Check some invalid combinations + raises(ValueError, lambda: PinJoint('J', P, C, joint_axis=P.x + C.y)) + raises(ValueError, lambda: PinJoint( + 'J', P, C, parent_interframe=Pint, child_interframe=Cint, + joint_axis=Pint.x + C.y)) + raises(ValueError, lambda: PinJoint( + 'J', P, C, parent_interframe=Pint, child_interframe=Cint, + joint_axis=P.x + Cint.y)) + # Check valid special combination + N, A, P, C, Pint, Cint = _generate_body(True) + PinJoint('J', P, C, parent_interframe=Pint, child_interframe=Cint, + joint_axis=Pint.x + P.y) + # Check invalid zero vector + raises(Exception, lambda: PinJoint( + 'J', P, C, parent_interframe=Pint, child_interframe=Cint, + joint_axis=Vector(0))) + raises(Exception, lambda: PinJoint( + 'J', P, C, parent_interframe=Pint, child_interframe=Cint, + joint_axis=P.y + Pint.y)) + + +def test_pin_joint_arbitrary_axis(): + q, u = dynamicsymbols('q_J, u_J') + + # When the bodies are attached though masscenters but axes are opposite. + N, A, P, C = _generate_body() + PinJoint('J', P, C, child_interframe=-A.x) + + assert (-A.x).angle_between(N.x) == 0 + assert -A.x.express(N) == N.x + assert A.dcm(N) == Matrix([[-1, 0, 0], + [0, -cos(q), -sin(q)], + [0, -sin(q), cos(q)]]) + assert A.ang_vel_in(N) == u*N.x + assert A.ang_vel_in(N).magnitude() == sqrt(u**2) + assert C.masscenter.pos_from(P.masscenter) == 0 + assert C.masscenter.pos_from(P.masscenter).express(N).simplify() == 0 + assert C.masscenter.vel(N) == 0 + + # When axes are different and parent joint is at masscenter but child joint + # is at a unit vector from child masscenter. + N, A, P, C = _generate_body() + PinJoint('J', P, C, child_interframe=A.y, child_point=A.x) + + assert A.y.angle_between(N.x) == 0 # Axis are aligned + assert A.y.express(N) == N.x + assert A.dcm(N) == Matrix([[0, -cos(q), -sin(q)], + [1, 0, 0], + [0, -sin(q), cos(q)]]) + assert A.ang_vel_in(N) == u*N.x + assert A.ang_vel_in(N).express(A) == u * A.y + assert A.ang_vel_in(N).magnitude() == sqrt(u**2) + assert A.ang_vel_in(N).cross(A.y) == 0 + assert C.masscenter.vel(N) == u*A.z + assert C.masscenter.pos_from(P.masscenter) == -A.x + assert (C.masscenter.pos_from(P.masscenter).express(N).simplify() == + cos(q)*N.y + sin(q)*N.z) + assert C.masscenter.vel(N).angle_between(A.x) == pi/2 + + # Similar to previous case but wrt parent body + N, A, P, C = _generate_body() + PinJoint('J', P, C, parent_interframe=N.y, parent_point=N.x) + + assert N.y.angle_between(A.x) == 0 # Axis are aligned + assert N.y.express(A) == A.x + assert A.dcm(N) == Matrix([[0, 1, 0], + [-cos(q), 0, sin(q)], + [sin(q), 0, cos(q)]]) + assert A.ang_vel_in(N) == u*N.y + assert A.ang_vel_in(N).express(A) == u*A.x + assert A.ang_vel_in(N).magnitude() == sqrt(u**2) + angle = A.ang_vel_in(N).angle_between(A.x) + assert angle.xreplace({u: 1}) == 0 + assert C.masscenter.vel(N) == 0 + assert C.masscenter.pos_from(P.masscenter) == N.x + + # Both joint pos id defined but different axes + N, A, P, C = _generate_body() + PinJoint('J', P, C, parent_point=N.x, child_point=A.x, + child_interframe=A.x + A.y) + assert expand_mul(N.x.angle_between(A.x + A.y)) == 0 # Axis are aligned + assert (A.x + A.y).express(N).simplify() == sqrt(2)*N.x + assert simplify(A.dcm(N)) == Matrix([ + [sqrt(2)/2, -sqrt(2)*cos(q)/2, -sqrt(2)*sin(q)/2], + [sqrt(2)/2, sqrt(2)*cos(q)/2, sqrt(2)*sin(q)/2], + [0, -sin(q), cos(q)]]) + assert A.ang_vel_in(N) == u*N.x + assert (A.ang_vel_in(N).express(A).simplify() == + (u*A.x + u*A.y)/sqrt(2)) + assert A.ang_vel_in(N).magnitude() == sqrt(u**2) + angle = A.ang_vel_in(N).angle_between(A.x + A.y) + assert angle.xreplace({u: 1}) == 0 + assert C.masscenter.vel(N).simplify() == (u * A.z)/sqrt(2) + assert C.masscenter.pos_from(P.masscenter) == N.x - A.x + assert (C.masscenter.pos_from(P.masscenter).express(N).simplify() == + (1 - sqrt(2)/2)*N.x + sqrt(2)*cos(q)/2*N.y + + sqrt(2)*sin(q)/2*N.z) + assert (C.masscenter.vel(N).express(N).simplify() == + -sqrt(2)*u*sin(q)/2*N.y + sqrt(2)*u*cos(q)/2*N.z) + assert C.masscenter.vel(N).angle_between(A.x) == pi/2 + + N, A, P, C = _generate_body() + PinJoint('J', P, C, parent_point=N.x, child_point=A.x, + child_interframe=A.x + A.y - A.z) + assert expand_mul(N.x.angle_between(A.x + A.y - A.z)) == 0 # Axis aligned + assert (A.x + A.y - A.z).express(N).simplify() == sqrt(3)*N.x + assert simplify(A.dcm(N)) == Matrix([ + [sqrt(3)/3, -sqrt(6)*sin(q + pi/4)/3, + sqrt(6)*cos(q + pi/4)/3], + [sqrt(3)/3, sqrt(6)*cos(q + pi/12)/3, + sqrt(6)*sin(q + pi/12)/3], + [-sqrt(3)/3, sqrt(6)*cos(q + 5*pi/12)/3, + sqrt(6)*sin(q + 5*pi/12)/3]]) + assert A.ang_vel_in(N) == u*N.x + assert A.ang_vel_in(N).express(A).simplify() == (u*A.x + u*A.y - + u*A.z)/sqrt(3) + assert A.ang_vel_in(N).magnitude() == sqrt(u**2) + angle = A.ang_vel_in(N).angle_between(A.x + A.y-A.z) + assert angle.xreplace({u: 1}).simplify() == 0 + assert C.masscenter.vel(N).simplify() == (u*A.y + u*A.z)/sqrt(3) + assert C.masscenter.pos_from(P.masscenter) == N.x - A.x + assert (C.masscenter.pos_from(P.masscenter).express(N).simplify() == + (1 - sqrt(3)/3)*N.x + sqrt(6)*sin(q + pi/4)/3*N.y - + sqrt(6)*cos(q + pi/4)/3*N.z) + assert (C.masscenter.vel(N).express(N).simplify() == + sqrt(6)*u*cos(q + pi/4)/3*N.y + + sqrt(6)*u*sin(q + pi/4)/3*N.z) + assert C.masscenter.vel(N).angle_between(A.x) == pi/2 + + N, A, P, C = _generate_body() + m, n = symbols('m n') + PinJoint('J', P, C, parent_point=m * N.x, child_point=n * A.x, + child_interframe=A.x + A.y - A.z, + parent_interframe=N.x - N.y + N.z) + angle = (N.x - N.y + N.z).angle_between(A.x + A.y - A.z) + assert expand_mul(angle) == 0 # Axis are aligned + assert ((A.x-A.y+A.z).express(N).simplify() == + (-4*cos(q)/3 - S(1)/3)*N.x + (S(1)/3 - 4*sin(q + pi/6)/3)*N.y + + (4*cos(q + pi/3)/3 - S(1)/3)*N.z) + assert simplify(A.dcm(N)) == Matrix([ + [S(1)/3 - 2*cos(q)/3, -2*sin(q + pi/6)/3 - S(1)/3, + 2*cos(q + pi/3)/3 + S(1)/3], + [2*cos(q + pi/3)/3 + S(1)/3, 2*cos(q)/3 - S(1)/3, + 2*sin(q + pi/6)/3 + S(1)/3], + [-2*sin(q + pi/6)/3 - S(1)/3, 2*cos(q + pi/3)/3 + S(1)/3, + 2*cos(q)/3 - S(1)/3]]) + assert (A.ang_vel_in(N) - (u*N.x - u*N.y + u*N.z)/sqrt(3)).simplify() + assert A.ang_vel_in(N).express(A).simplify() == (u*A.x + u*A.y - + u*A.z)/sqrt(3) + assert A.ang_vel_in(N).magnitude() == sqrt(u**2) + angle = A.ang_vel_in(N).angle_between(A.x+A.y-A.z) + assert angle.xreplace({u: 1}).simplify() == 0 + assert (C.masscenter.vel(N).simplify() == + sqrt(3)*n*u/3*A.y + sqrt(3)*n*u/3*A.z) + assert C.masscenter.pos_from(P.masscenter) == m*N.x - n*A.x + assert (C.masscenter.pos_from(P.masscenter).express(N).simplify() == + (m + n*(2*cos(q) - 1)/3)*N.x + n*(2*sin(q + pi/6) + + 1)/3*N.y - n*(2*cos(q + pi/3) + 1)/3*N.z) + assert (C.masscenter.vel(N).express(N).simplify() == + - 2*n*u*sin(q)/3*N.x + 2*n*u*cos(q + pi/6)/3*N.y + + 2*n*u*sin(q + pi/3)/3*N.z) + assert C.masscenter.vel(N).dot(N.x - N.y + N.z).simplify() == 0 + + +def test_create_aligned_frame_pi(): + N, A, P, C = _generate_body() + f = Joint._create_aligned_interframe(P, -P.x, P.x) + assert f.z == P.z + f = Joint._create_aligned_interframe(P, -P.y, P.y) + assert f.x == P.x + f = Joint._create_aligned_interframe(P, -P.z, P.z) + assert f.y == P.y + f = Joint._create_aligned_interframe(P, -P.x - P.y, P.x + P.y) + assert f.z == P.z + f = Joint._create_aligned_interframe(P, -P.y - P.z, P.y + P.z) + assert f.x == P.x + f = Joint._create_aligned_interframe(P, -P.x - P.z, P.x + P.z) + assert f.y == P.y + f = Joint._create_aligned_interframe(P, -P.x - P.y - P.z, P.x + P.y + P.z) + assert f.y - f.z == P.y - P.z + + +def test_pin_joint_axis(): + q, u = dynamicsymbols('q u') + # Test default joint axis + N, A, P, C, Pint, Cint = _generate_body(True) + J = PinJoint('J', P, C, q, u, parent_interframe=Pint, child_interframe=Cint) + assert J.joint_axis == Pint.x + # Test for the same joint axis expressed in different frames + N_R_A = Matrix([[0, sin(q), cos(q)], + [0, -cos(q), sin(q)], + [1, 0, 0]]) + N, A, P, C, Pint, Cint = _generate_body(True) + PinJoint('J', P, C, q, u, parent_interframe=Pint, child_interframe=Cint, + joint_axis=N.z) + assert N.dcm(A) == N_R_A + N, A, P, C, Pint, Cint = _generate_body(True) + PinJoint('J', P, C, q, u, parent_interframe=Pint, child_interframe=Cint, + joint_axis=-Pint.z) + assert N.dcm(A) == N_R_A + # Test time varying joint axis + N, A, P, C, Pint, Cint = _generate_body(True) + raises(ValueError, lambda: PinJoint('J', P, C, joint_axis=q * N.z)) + + +def test_locate_joint_pos(): + # Test Vector and default + N, A, P, C = _generate_body() + joint = PinJoint('J', P, C, parent_point=N.y + N.z) + assert joint.parent_point.name == 'J_P_joint' + assert joint.parent_point.pos_from(P.masscenter) == N.y + N.z + assert joint.child_point == C.masscenter + # Test Point objects + N, A, P, C = _generate_body() + parent_point = P.masscenter.locatenew('p', N.y + N.z) + joint = PinJoint('J', P, C, parent_point=parent_point, + child_point=C.masscenter) + assert joint.parent_point == parent_point + assert joint.child_point == C.masscenter + # Check invalid type + N, A, P, C = _generate_body() + raises(TypeError, + lambda: PinJoint('J', P, C, parent_point=N.x.to_matrix(N))) + # Test time varying positions + q = dynamicsymbols('q') + N, A, P, C = _generate_body() + raises(ValueError, lambda: PinJoint('J', P, C, parent_point=q * N.x)) + N, A, P, C = _generate_body() + child_point = C.masscenter.locatenew('p', q * A.y) + raises(ValueError, lambda: PinJoint('J', P, C, child_point=child_point)) + # Test undefined position + child_point = Point('p') + raises(ValueError, lambda: PinJoint('J', P, C, child_point=child_point)) + + +def test_locate_joint_frame(): + # Test rotated frame and default + N, A, P, C = _generate_body() + parent_interframe = ReferenceFrame('int_frame') + parent_interframe.orient_axis(N, N.z, 1) + joint = PinJoint('J', P, C, parent_interframe=parent_interframe) + assert joint.parent_interframe == parent_interframe + assert joint.parent_interframe.ang_vel_in(N) == 0 + assert joint.child_interframe == A + # Test time varying orientations + q = dynamicsymbols('q') + N, A, P, C = _generate_body() + parent_interframe = ReferenceFrame('int_frame') + parent_interframe.orient_axis(N, N.z, q) + raises(ValueError, + lambda: PinJoint('J', P, C, parent_interframe=parent_interframe)) + # Test undefined frame + N, A, P, C = _generate_body() + child_interframe = ReferenceFrame('int_frame') + child_interframe.orient_axis(N, N.z, 1) # Defined with respect to parent + raises(ValueError, + lambda: PinJoint('J', P, C, child_interframe=child_interframe)) + + +def test_prismatic_joint(): + _, _, P, C = _generate_body() + q, u = dynamicsymbols('q_S, u_S') + S = PrismaticJoint('S', P, C) + assert S.name == 'S' + assert S.parent == P + assert S.child == C + assert S.coordinates == Matrix([q]) + assert S.speeds == Matrix([u]) + assert S.kdes == Matrix([u - q.diff(t)]) + assert S.joint_axis == P.frame.x + assert S.child_point.pos_from(C.masscenter) == Vector(0) + assert S.parent_point.pos_from(P.masscenter) == Vector(0) + assert S.parent_point.pos_from(S.child_point) == - q * P.frame.x + assert P.masscenter.pos_from(C.masscenter) == - q * P.frame.x + assert C.masscenter.vel(P.frame) == u * P.frame.x + assert P.frame.ang_vel_in(C.frame) == 0 + assert C.frame.ang_vel_in(P.frame) == 0 + assert S.__str__() == 'PrismaticJoint: S parent: P child: C' + + N, A, P, C = _generate_body() + l, m = symbols('l m') + Pint = ReferenceFrame('P_int') + Pint.orient_axis(P.frame, P.y, pi / 2) + S = PrismaticJoint('S', P, C, parent_point=l * P.frame.x, + child_point=m * C.frame.y, joint_axis=P.frame.z, + parent_interframe=Pint) + + assert S.joint_axis == P.frame.z + assert S.child_point.pos_from(C.masscenter) == m * C.frame.y + assert S.parent_point.pos_from(P.masscenter) == l * P.frame.x + assert S.parent_point.pos_from(S.child_point) == - q * P.frame.z + assert P.masscenter.pos_from(C.masscenter) == - l * N.x - q * N.z + m * A.y + assert C.masscenter.vel(P.frame) == u * P.frame.z + assert P.masscenter.vel(Pint) == Vector(0) + assert C.frame.ang_vel_in(P.frame) == 0 + assert P.frame.ang_vel_in(C.frame) == 0 + + _, _, P, C = _generate_body() + Pint = ReferenceFrame('P_int') + Pint.orient_axis(P.frame, P.y, pi / 2) + S = PrismaticJoint('S', P, C, parent_point=l * P.frame.z, + child_point=m * C.frame.x, joint_axis=P.frame.z, + parent_interframe=Pint) + assert S.joint_axis == P.frame.z + assert S.child_point.pos_from(C.masscenter) == m * C.frame.x + assert S.parent_point.pos_from(P.masscenter) == l * P.frame.z + assert S.parent_point.pos_from(S.child_point) == - q * P.frame.z + assert P.masscenter.pos_from(C.masscenter) == (-l - q)*P.frame.z + m*C.frame.x + assert C.masscenter.vel(P.frame) == u * P.frame.z + assert C.frame.ang_vel_in(P.frame) == 0 + assert P.frame.ang_vel_in(C.frame) == 0 + + +def test_prismatic_joint_arbitrary_axis(): + q, u = dynamicsymbols('q_S, u_S') + + N, A, P, C = _generate_body() + PrismaticJoint('S', P, C, child_interframe=-A.x) + + assert (-A.x).angle_between(N.x) == 0 + assert -A.x.express(N) == N.x + assert A.dcm(N) == Matrix([[-1, 0, 0], [0, -1, 0], [0, 0, 1]]) + assert C.masscenter.pos_from(P.masscenter) == q * N.x + assert C.masscenter.pos_from(P.masscenter).express(A).simplify() == -q * A.x + assert C.masscenter.vel(N) == u * N.x + assert C.masscenter.vel(N).express(A) == -u * A.x + assert A.ang_vel_in(N) == 0 + assert N.ang_vel_in(A) == 0 + + #When axes are different and parent joint is at masscenter but child joint is at a unit vector from + #child masscenter. + N, A, P, C = _generate_body() + PrismaticJoint('S', P, C, child_interframe=A.y, child_point=A.x) + + assert A.y.angle_between(N.x) == 0 #Axis are aligned + assert A.y.express(N) == N.x + assert A.dcm(N) == Matrix([[0, -1, 0], [1, 0, 0], [0, 0, 1]]) + assert C.masscenter.vel(N) == u * N.x + assert C.masscenter.vel(N).express(A) == u * A.y + assert C.masscenter.pos_from(P.masscenter) == q*N.x - A.x + assert C.masscenter.pos_from(P.masscenter).express(N).simplify() == q*N.x + N.y + assert A.ang_vel_in(N) == 0 + assert N.ang_vel_in(A) == 0 + + #Similar to previous case but wrt parent body + N, A, P, C = _generate_body() + PrismaticJoint('S', P, C, parent_interframe=N.y, parent_point=N.x) + + assert N.y.angle_between(A.x) == 0 #Axis are aligned + assert N.y.express(A) == A.x + assert A.dcm(N) == Matrix([[0, 1, 0], [-1, 0, 0], [0, 0, 1]]) + assert C.masscenter.vel(N) == u * N.y + assert C.masscenter.vel(N).express(A) == u * A.x + assert C.masscenter.pos_from(P.masscenter) == N.x + q*N.y + assert A.ang_vel_in(N) == 0 + assert N.ang_vel_in(A) == 0 + + #Both joint pos is defined but different axes + N, A, P, C = _generate_body() + PrismaticJoint('S', P, C, parent_point=N.x, child_point=A.x, + child_interframe=A.x + A.y) + assert N.x.angle_between(A.x + A.y) == 0 #Axis are aligned + assert (A.x + A.y).express(N) == sqrt(2)*N.x + assert A.dcm(N) == Matrix([[sqrt(2)/2, -sqrt(2)/2, 0], [sqrt(2)/2, sqrt(2)/2, 0], [0, 0, 1]]) + assert C.masscenter.pos_from(P.masscenter) == (q + 1)*N.x - A.x + assert C.masscenter.pos_from(P.masscenter).express(N) == (q - sqrt(2)/2 + 1)*N.x + sqrt(2)/2*N.y + assert C.masscenter.vel(N).express(A) == u * (A.x + A.y)/sqrt(2) + assert C.masscenter.vel(N) == u*N.x + assert A.ang_vel_in(N) == 0 + assert N.ang_vel_in(A) == 0 + + N, A, P, C = _generate_body() + PrismaticJoint('S', P, C, parent_point=N.x, child_point=A.x, + child_interframe=A.x + A.y - A.z) + assert N.x.angle_between(A.x + A.y - A.z).simplify() == 0 #Axis are aligned + assert ((A.x + A.y - A.z).express(N) - sqrt(3)*N.x).simplify() == 0 + assert simplify(A.dcm(N)) == Matrix([[sqrt(3)/3, -sqrt(3)/3, sqrt(3)/3], + [sqrt(3)/3, sqrt(3)/6 + S(1)/2, S(1)/2 - sqrt(3)/6], + [-sqrt(3)/3, S(1)/2 - sqrt(3)/6, sqrt(3)/6 + S(1)/2]]) + assert C.masscenter.pos_from(P.masscenter) == (q + 1)*N.x - A.x + assert (C.masscenter.pos_from(P.masscenter).express(N) - + ((q - sqrt(3)/3 + 1)*N.x + sqrt(3)/3*N.y - sqrt(3)/3*N.z)).simplify() == 0 + assert C.masscenter.vel(N) == u*N.x + assert (C.masscenter.vel(N).express(A) - ( + sqrt(3)*u/3*A.x + sqrt(3)*u/3*A.y - sqrt(3)*u/3*A.z)).simplify() + assert A.ang_vel_in(N) == 0 + assert N.ang_vel_in(A) == 0 + + N, A, P, C = _generate_body() + m, n = symbols('m n') + PrismaticJoint('S', P, C, parent_point=m*N.x, child_point=n*A.x, + child_interframe=A.x + A.y - A.z, + parent_interframe=N.x - N.y + N.z) + # 0 angle means that the axis are aligned + assert (N.x-N.y+N.z).angle_between(A.x+A.y-A.z).simplify() == 0 + assert ((A.x+A.y-A.z).express(N) - (N.x - N.y + N.z)).simplify() == 0 + assert simplify(A.dcm(N)) == Matrix([[-S(1)/3, -S(2)/3, S(2)/3], + [S(2)/3, S(1)/3, S(2)/3], + [-S(2)/3, S(2)/3, S(1)/3]]) + assert (C.masscenter.pos_from(P.masscenter) - ( + (m + sqrt(3)*q/3)*N.x - sqrt(3)*q/3*N.y + sqrt(3)*q/3*N.z - n*A.x) + ).express(N).simplify() == 0 + assert (C.masscenter.pos_from(P.masscenter).express(N) - ( + (m + n/3 + sqrt(3)*q/3)*N.x + (2*n/3 - sqrt(3)*q/3)*N.y + + (-2*n/3 + sqrt(3)*q/3)*N.z)).simplify() == 0 + assert (C.masscenter.vel(N).express(N) - ( + sqrt(3)*u/3*N.x - sqrt(3)*u/3*N.y + sqrt(3)*u/3*N.z)).simplify() == 0 + assert (C.masscenter.vel(N).express(A) - + (sqrt(3)*u/3*A.x + sqrt(3)*u/3*A.y - sqrt(3)*u/3*A.z)).simplify() == 0 + assert A.ang_vel_in(N) == 0 + assert N.ang_vel_in(A) == 0 + + +def test_cylindrical_joint(): + N, A, P, C = _generate_body() + q0_def, q1_def, u0_def, u1_def = dynamicsymbols('q0:2_J, u0:2_J') + Cj = CylindricalJoint('J', P, C) + assert Cj.name == 'J' + assert Cj.parent == P + assert Cj.child == C + assert Cj.coordinates == Matrix([q0_def, q1_def]) + assert Cj.speeds == Matrix([u0_def, u1_def]) + assert Cj.rotation_coordinate == q0_def + assert Cj.translation_coordinate == q1_def + assert Cj.rotation_speed == u0_def + assert Cj.translation_speed == u1_def + assert Cj.kdes == Matrix([u0_def - q0_def.diff(t), u1_def - q1_def.diff(t)]) + assert Cj.joint_axis == N.x + assert Cj.child_point.pos_from(C.masscenter) == Vector(0) + assert Cj.parent_point.pos_from(P.masscenter) == Vector(0) + assert Cj.parent_point.pos_from(Cj._child_point) == -q1_def * N.x + assert C.masscenter.pos_from(P.masscenter) == q1_def * N.x + assert Cj.child_point.vel(N) == u1_def * N.x + assert A.ang_vel_in(N) == u0_def * N.x + assert Cj.parent_interframe == N + assert Cj.child_interframe == A + assert Cj.__str__() == 'CylindricalJoint: J parent: P child: C' + + q0, q1, u0, u1 = dynamicsymbols('q0:2, u0:2') + l, m = symbols('l, m') + N, A, P, C, Pint, Cint = _generate_body(True) + Cj = CylindricalJoint('J', P, C, rotation_coordinate=q0, rotation_speed=u0, + translation_speed=u1, parent_point=m * N.x, + child_point=l * A.y, parent_interframe=Pint, + child_interframe=Cint, joint_axis=2 * N.z) + assert Cj.coordinates == Matrix([q0, q1_def]) + assert Cj.speeds == Matrix([u0, u1]) + assert Cj.rotation_coordinate == q0 + assert Cj.translation_coordinate == q1_def + assert Cj.rotation_speed == u0 + assert Cj.translation_speed == u1 + assert Cj.kdes == Matrix([u0 - q0.diff(t), u1 - q1_def.diff(t)]) + assert Cj.joint_axis == 2 * N.z + assert Cj.child_point.pos_from(C.masscenter) == l * A.y + assert Cj.parent_point.pos_from(P.masscenter) == m * N.x + assert Cj.parent_point.pos_from(Cj._child_point) == -q1_def * N.z + assert C.masscenter.pos_from( + P.masscenter) == m * N.x + q1_def * N.z - l * A.y + assert C.masscenter.vel(N) == u1 * N.z - u0 * l * A.z + assert A.ang_vel_in(N) == u0 * N.z + + +def test_planar_joint(): + N, A, P, C = _generate_body() + q0_def, q1_def, q2_def = dynamicsymbols('q0:3_J') + u0_def, u1_def, u2_def = dynamicsymbols('u0:3_J') + Cj = PlanarJoint('J', P, C) + assert Cj.name == 'J' + assert Cj.parent == P + assert Cj.child == C + assert Cj.coordinates == Matrix([q0_def, q1_def, q2_def]) + assert Cj.speeds == Matrix([u0_def, u1_def, u2_def]) + assert Cj.rotation_coordinate == q0_def + assert Cj.planar_coordinates == Matrix([q1_def, q2_def]) + assert Cj.rotation_speed == u0_def + assert Cj.planar_speeds == Matrix([u1_def, u2_def]) + assert Cj.kdes == Matrix([u0_def - q0_def.diff(t), u1_def - q1_def.diff(t), + u2_def - q2_def.diff(t)]) + assert Cj.rotation_axis == N.x + assert Cj.planar_vectors == [N.y, N.z] + assert Cj.child_point.pos_from(C.masscenter) == Vector(0) + assert Cj.parent_point.pos_from(P.masscenter) == Vector(0) + r_P_C = q1_def * N.y + q2_def * N.z + assert Cj.parent_point.pos_from(Cj.child_point) == -r_P_C + assert C.masscenter.pos_from(P.masscenter) == r_P_C + assert Cj.child_point.vel(N) == u1_def * N.y + u2_def * N.z + assert A.ang_vel_in(N) == u0_def * N.x + assert Cj.parent_interframe == N + assert Cj.child_interframe == A + assert Cj.__str__() == 'PlanarJoint: J parent: P child: C' + + q0, q1, q2, u0, u1, u2 = dynamicsymbols('q0:3, u0:3') + l, m = symbols('l, m') + N, A, P, C, Pint, Cint = _generate_body(True) + Cj = PlanarJoint('J', P, C, rotation_coordinate=q0, + planar_coordinates=[q1, q2], planar_speeds=[u1, u2], + parent_point=m * N.x, child_point=l * A.y, + parent_interframe=Pint, child_interframe=Cint) + assert Cj.coordinates == Matrix([q0, q1, q2]) + assert Cj.speeds == Matrix([u0_def, u1, u2]) + assert Cj.rotation_coordinate == q0 + assert Cj.planar_coordinates == Matrix([q1, q2]) + assert Cj.rotation_speed == u0_def + assert Cj.planar_speeds == Matrix([u1, u2]) + assert Cj.kdes == Matrix([u0_def - q0.diff(t), u1 - q1.diff(t), + u2 - q2.diff(t)]) + assert Cj.rotation_axis == Pint.x + assert Cj.planar_vectors == [Pint.y, Pint.z] + assert Cj.child_point.pos_from(C.masscenter) == l * A.y + assert Cj.parent_point.pos_from(P.masscenter) == m * N.x + assert Cj.parent_point.pos_from(Cj.child_point) == q1 * N.y + q2 * N.z + assert C.masscenter.pos_from( + P.masscenter) == m * N.x - q1 * N.y - q2 * N.z - l * A.y + assert C.masscenter.vel(N) == -u1 * N.y - u2 * N.z + u0_def * l * A.x + assert A.ang_vel_in(N) == u0_def * N.x + + +def test_planar_joint_advanced(): + # Tests whether someone is able to just specify two normals, which will form + # the rotation axis seen from the parent and child body. + # This specific example is a block on a slope, which has that same slope of + # 30 degrees, so in the zero configuration the frames of the parent and + # child are actually aligned. + q0, q1, q2, u0, u1, u2 = dynamicsymbols('q0:3, u0:3') + l1, l2 = symbols('l1:3') + N, A, P, C = _generate_body() + J = PlanarJoint('J', P, C, q0, [q1, q2], u0, [u1, u2], + parent_point=l1 * N.z, + child_point=-l2 * C.z, + parent_interframe=N.z + N.y / sqrt(3), + child_interframe=A.z + A.y / sqrt(3)) + assert J.rotation_axis.express(N) == (N.z + N.y / sqrt(3)).normalize() + assert J.rotation_axis.express(A) == (A.z + A.y / sqrt(3)).normalize() + assert J.rotation_axis.angle_between(N.z) == pi / 6 + assert N.dcm(A).xreplace({q0: 0, q1: 0, q2: 0}) == eye(3) + N_R_A = Matrix([ + [cos(q0), -sqrt(3) * sin(q0) / 2, sin(q0) / 2], + [sqrt(3) * sin(q0) / 2, 3 * cos(q0) / 4 + 1 / 4, + sqrt(3) * (1 - cos(q0)) / 4], + [-sin(q0) / 2, sqrt(3) * (1 - cos(q0)) / 4, cos(q0) / 4 + 3 / 4]]) + # N.dcm(A) == N_R_A did not work + assert simplify(N.dcm(A) - N_R_A) == zeros(3) + + +def test_spherical_joint(): + N, A, P, C = _generate_body() + q0, q1, q2, u0, u1, u2 = dynamicsymbols('q0:3_S, u0:3_S') + S = SphericalJoint('S', P, C) + assert S.name == 'S' + assert S.parent == P + assert S.child == C + assert S.coordinates == Matrix([q0, q1, q2]) + assert S.speeds == Matrix([u0, u1, u2]) + assert S.kdes == Matrix([u0 - q0.diff(t), u1 - q1.diff(t), u2 - q2.diff(t)]) + assert S.child_point.pos_from(C.masscenter) == Vector(0) + assert S.parent_point.pos_from(P.masscenter) == Vector(0) + assert S.parent_point.pos_from(S.child_point) == Vector(0) + assert P.masscenter.pos_from(C.masscenter) == Vector(0) + assert C.masscenter.vel(N) == Vector(0) + assert N.ang_vel_in(A) == (-u0 * cos(q1) * cos(q2) - u1 * sin(q2)) * A.x + ( + u0 * sin(q2) * cos(q1) - u1 * cos(q2)) * A.y + ( + -u0 * sin(q1) - u2) * A.z + assert A.ang_vel_in(N) == (u0 * cos(q1) * cos(q2) + u1 * sin(q2)) * A.x + ( + -u0 * sin(q2) * cos(q1) + u1 * cos(q2)) * A.y + ( + u0 * sin(q1) + u2) * A.z + assert S.__str__() == 'SphericalJoint: S parent: P child: C' + assert S._rot_type == 'BODY' + assert S._rot_order == 123 + assert S._amounts is None + + +def test_spherical_joint_speeds_as_derivative_terms(): + # This tests checks whether the system remains valid if the user chooses to + # pass the derivative of the generalized coordinates as generalized speeds + q0, q1, q2 = dynamicsymbols('q0:3') + u0, u1, u2 = dynamicsymbols('q0:3', 1) + N, A, P, C = _generate_body() + S = SphericalJoint('S', P, C, coordinates=[q0, q1, q2], speeds=[u0, u1, u2]) + assert S.coordinates == Matrix([q0, q1, q2]) + assert S.speeds == Matrix([u0, u1, u2]) + assert S.kdes == Matrix([0, 0, 0]) + assert N.ang_vel_in(A) == (-u0 * cos(q1) * cos(q2) - u1 * sin(q2)) * A.x + ( + u0 * sin(q2) * cos(q1) - u1 * cos(q2)) * A.y + ( + -u0 * sin(q1) - u2) * A.z + + +def test_spherical_joint_coords(): + q0s, q1s, q2s, u0s, u1s, u2s = dynamicsymbols('q0:3_S, u0:3_S') + q0, q1, q2, q3, u0, u1, u2, u4 = dynamicsymbols('q0:4, u0:4') + # Test coordinates as list + N, A, P, C = _generate_body() + S = SphericalJoint('S', P, C, [q0, q1, q2], [u0, u1, u2]) + assert S.coordinates == Matrix([q0, q1, q2]) + assert S.speeds == Matrix([u0, u1, u2]) + # Test coordinates as Matrix + N, A, P, C = _generate_body() + S = SphericalJoint('S', P, C, Matrix([q0, q1, q2]), + Matrix([u0, u1, u2])) + assert S.coordinates == Matrix([q0, q1, q2]) + assert S.speeds == Matrix([u0, u1, u2]) + # Test too few generalized coordinates + N, A, P, C = _generate_body() + raises(ValueError, + lambda: SphericalJoint('S', P, C, Matrix([q0, q1]), Matrix([u0]))) + # Test too many generalized coordinates + raises(ValueError, lambda: SphericalJoint( + 'S', P, C, Matrix([q0, q1, q2, q3]), Matrix([u0, u1, u2]))) + raises(ValueError, lambda: SphericalJoint( + 'S', P, C, Matrix([q0, q1, q2]), Matrix([u0, u1, u2, u4]))) + + +def test_spherical_joint_orient_body(): + q0, q1, q2, u0, u1, u2 = dynamicsymbols('q0:3, u0:3') + N_R_A = Matrix([ + [-sin(q1), -sin(q2) * cos(q1), cos(q1) * cos(q2)], + [-sin(q0) * cos(q1), sin(q0) * sin(q1) * sin(q2) - cos(q0) * cos(q2), + -sin(q0) * sin(q1) * cos(q2) - sin(q2) * cos(q0)], + [cos(q0) * cos(q1), -sin(q0) * cos(q2) - sin(q1) * sin(q2) * cos(q0), + -sin(q0) * sin(q2) + sin(q1) * cos(q0) * cos(q2)]]) + N_w_A = Matrix([[-u0 * sin(q1) - u2], + [-u0 * sin(q2) * cos(q1) + u1 * cos(q2)], + [u0 * cos(q1) * cos(q2) + u1 * sin(q2)]]) + N_v_Co = Matrix([ + [-sqrt(2) * (u0 * cos(q2 + pi / 4) * cos(q1) + u1 * sin(q2 + pi / 4))], + [-u0 * sin(q1) - u2], [-u0 * sin(q1) - u2]]) + # Test default rot_type='BODY', rot_order=123 + N, A, P, C, Pint, Cint = _generate_body(True) + S = SphericalJoint('S', P, C, coordinates=[q0, q1, q2], speeds=[u0, u1, u2], + parent_point=N.x + N.y, child_point=-A.y + A.z, + parent_interframe=Pint, child_interframe=Cint, + rot_type='body', rot_order=123) + assert S._rot_type.upper() == 'BODY' + assert S._rot_order == 123 + assert simplify(N.dcm(A) - N_R_A) == zeros(3) + assert simplify(A.ang_vel_in(N).to_matrix(A) - N_w_A) == zeros(3, 1) + assert simplify(C.masscenter.vel(N).to_matrix(A)) == N_v_Co + # Test change of amounts + N, A, P, C, Pint, Cint = _generate_body(True) + S = SphericalJoint('S', P, C, coordinates=[q0, q1, q2], speeds=[u0, u1, u2], + parent_point=N.x + N.y, child_point=-A.y + A.z, + parent_interframe=Pint, child_interframe=Cint, + rot_type='BODY', amounts=(q1, q0, q2), rot_order=123) + switch_order = lambda expr: expr.xreplace( + {q0: q1, q1: q0, q2: q2, u0: u1, u1: u0, u2: u2}) + assert S._rot_type.upper() == 'BODY' + assert S._rot_order == 123 + assert simplify(N.dcm(A) - switch_order(N_R_A)) == zeros(3) + assert simplify(A.ang_vel_in(N).to_matrix(A) - switch_order(N_w_A) + ) == zeros(3, 1) + assert simplify(C.masscenter.vel(N).to_matrix(A)) == switch_order(N_v_Co) + # Test different rot_order + N, A, P, C, Pint, Cint = _generate_body(True) + S = SphericalJoint('S', P, C, coordinates=[q0, q1, q2], speeds=[u0, u1, u2], + parent_point=N.x + N.y, child_point=-A.y + A.z, + parent_interframe=Pint, child_interframe=Cint, + rot_type='BodY', rot_order='yxz') + assert S._rot_type.upper() == 'BODY' + assert S._rot_order == 'yxz' + assert simplify(N.dcm(A) - Matrix([ + [-sin(q0) * cos(q1), sin(q0) * sin(q1) * cos(q2) - sin(q2) * cos(q0), + sin(q0) * sin(q1) * sin(q2) + cos(q0) * cos(q2)], + [-sin(q1), -cos(q1) * cos(q2), -sin(q2) * cos(q1)], + [cos(q0) * cos(q1), -sin(q0) * sin(q2) - sin(q1) * cos(q0) * cos(q2), + sin(q0) * cos(q2) - sin(q1) * sin(q2) * cos(q0)]])) == zeros(3) + assert simplify(A.ang_vel_in(N).to_matrix(A) - Matrix([ + [u0 * sin(q1) - u2], [u0 * cos(q1) * cos(q2) - u1 * sin(q2)], + [u0 * sin(q2) * cos(q1) + u1 * cos(q2)]])) == zeros(3, 1) + assert simplify(C.masscenter.vel(N).to_matrix(A)) == Matrix([ + [-sqrt(2) * (u0 * sin(q2 + pi / 4) * cos(q1) + u1 * cos(q2 + pi / 4))], + [u0 * sin(q1) - u2], [u0 * sin(q1) - u2]]) + + +def test_spherical_joint_orient_space(): + q0, q1, q2, u0, u1, u2 = dynamicsymbols('q0:3, u0:3') + N_R_A = Matrix([ + [-sin(q0) * sin(q2) - sin(q1) * cos(q0) * cos(q2), + sin(q0) * sin(q1) * cos(q2) - sin(q2) * cos(q0), cos(q1) * cos(q2)], + [-sin(q0) * cos(q2) + sin(q1) * sin(q2) * cos(q0), + -sin(q0) * sin(q1) * sin(q2) - cos(q0) * cos(q2), -sin(q2) * cos(q1)], + [cos(q0) * cos(q1), -sin(q0) * cos(q1), sin(q1)]]) + N_w_A = Matrix([ + [u1 * sin(q0) - u2 * cos(q0) * cos(q1)], + [u1 * cos(q0) + u2 * sin(q0) * cos(q1)], [u0 - u2 * sin(q1)]]) + N_v_Co = Matrix([ + [u0 - u2 * sin(q1)], [u0 - u2 * sin(q1)], + [sqrt(2) * (-u1 * sin(q0 + pi / 4) + u2 * cos(q0 + pi / 4) * cos(q1))]]) + # Test default rot_type='BODY', rot_order=123 + N, A, P, C, Pint, Cint = _generate_body(True) + S = SphericalJoint('S', P, C, coordinates=[q0, q1, q2], speeds=[u0, u1, u2], + parent_point=N.x + N.z, child_point=-A.x + A.y, + parent_interframe=Pint, child_interframe=Cint, + rot_type='space', rot_order=123) + assert S._rot_type.upper() == 'SPACE' + assert S._rot_order == 123 + assert simplify(N.dcm(A) - N_R_A) == zeros(3) + assert simplify(A.ang_vel_in(N).to_matrix(A)) == N_w_A + assert simplify(C.masscenter.vel(N).to_matrix(A)) == N_v_Co + # Test change of amounts + switch_order = lambda expr: expr.xreplace( + {q0: q1, q1: q0, q2: q2, u0: u1, u1: u0, u2: u2}) + N, A, P, C, Pint, Cint = _generate_body(True) + S = SphericalJoint('S', P, C, coordinates=[q0, q1, q2], speeds=[u0, u1, u2], + parent_point=N.x + N.z, child_point=-A.x + A.y, + parent_interframe=Pint, child_interframe=Cint, + rot_type='SPACE', amounts=(q1, q0, q2), rot_order=123) + assert S._rot_type.upper() == 'SPACE' + assert S._rot_order == 123 + assert simplify(N.dcm(A) - switch_order(N_R_A)) == zeros(3) + assert simplify(A.ang_vel_in(N).to_matrix(A)) == switch_order(N_w_A) + assert simplify(C.masscenter.vel(N).to_matrix(A)) == switch_order(N_v_Co) + # Test different rot_order + N, A, P, C, Pint, Cint = _generate_body(True) + S = SphericalJoint('S', P, C, coordinates=[q0, q1, q2], speeds=[u0, u1, u2], + parent_point=N.x + N.z, child_point=-A.x + A.y, + parent_interframe=Pint, child_interframe=Cint, + rot_type='SPaCe', rot_order='zxy') + assert S._rot_type.upper() == 'SPACE' + assert S._rot_order == 'zxy' + assert simplify(N.dcm(A) - Matrix([ + [-sin(q2) * cos(q1), -sin(q0) * cos(q2) + sin(q1) * sin(q2) * cos(q0), + sin(q0) * sin(q1) * sin(q2) + cos(q0) * cos(q2)], + [-sin(q1), -cos(q0) * cos(q1), -sin(q0) * cos(q1)], + [cos(q1) * cos(q2), -sin(q0) * sin(q2) - sin(q1) * cos(q0) * cos(q2), + -sin(q0) * sin(q1) * cos(q2) + sin(q2) * cos(q0)]])) + assert simplify(A.ang_vel_in(N).to_matrix(A) - Matrix([ + [-u0 + u2 * sin(q1)], [-u1 * sin(q0) + u2 * cos(q0) * cos(q1)], + [u1 * cos(q0) + u2 * sin(q0) * cos(q1)]])) == zeros(3, 1) + assert simplify(C.masscenter.vel(N).to_matrix(A) - Matrix([ + [u1 * cos(q0) + u2 * sin(q0) * cos(q1)], + [u1 * cos(q0) + u2 * sin(q0) * cos(q1)], + [u0 + u1 * sin(q0) - u2 * sin(q1) - + u2 * cos(q0) * cos(q1)]])) == zeros(3, 1) + + +def test_weld_joint(): + _, _, P, C = _generate_body() + W = WeldJoint('W', P, C) + assert W.name == 'W' + assert W.parent == P + assert W.child == C + assert W.coordinates == Matrix() + assert W.speeds == Matrix() + assert W.kdes == Matrix(1, 0, []).T + assert P.frame.dcm(C.frame) == eye(3) + assert W.child_point.pos_from(C.masscenter) == Vector(0) + assert W.parent_point.pos_from(P.masscenter) == Vector(0) + assert W.parent_point.pos_from(W.child_point) == Vector(0) + assert P.masscenter.pos_from(C.masscenter) == Vector(0) + assert C.masscenter.vel(P.frame) == Vector(0) + assert P.frame.ang_vel_in(C.frame) == 0 + assert C.frame.ang_vel_in(P.frame) == 0 + assert W.__str__() == 'WeldJoint: W parent: P child: C' + + N, A, P, C = _generate_body() + l, m = symbols('l m') + Pint = ReferenceFrame('P_int') + Pint.orient_axis(P.frame, P.y, pi / 2) + W = WeldJoint('W', P, C, parent_point=l * P.frame.x, + child_point=m * C.frame.y, parent_interframe=Pint) + + assert W.child_point.pos_from(C.masscenter) == m * C.frame.y + assert W.parent_point.pos_from(P.masscenter) == l * P.frame.x + assert W.parent_point.pos_from(W.child_point) == Vector(0) + assert P.masscenter.pos_from(C.masscenter) == - l * N.x + m * A.y + assert C.masscenter.vel(P.frame) == Vector(0) + assert P.masscenter.vel(Pint) == Vector(0) + assert C.frame.ang_vel_in(P.frame) == 0 + assert P.frame.ang_vel_in(C.frame) == 0 + assert P.x == A.z + + with warns_deprecated_sympy(): + JointsMethod(P, W) # Tests #10770 + + +def test_deprecated_parent_child_axis(): + q, u = dynamicsymbols('q_J, u_J') + N, A, P, C = _generate_body() + with warns_deprecated_sympy(): + PinJoint('J', P, C, child_axis=-A.x) + assert (-A.x).angle_between(N.x) == 0 + assert -A.x.express(N) == N.x + assert A.dcm(N) == Matrix([[-1, 0, 0], + [0, -cos(q), -sin(q)], + [0, -sin(q), cos(q)]]) + assert A.ang_vel_in(N) == u * N.x + assert A.ang_vel_in(N).magnitude() == sqrt(u ** 2) + + N, A, P, C = _generate_body() + with warns_deprecated_sympy(): + PrismaticJoint('J', P, C, parent_axis=P.x + P.y) + assert (A.x).angle_between(N.x + N.y) == 0 + assert A.x.express(N) == (N.x + N.y) / sqrt(2) + assert A.dcm(N) == Matrix([[sqrt(2) / 2, sqrt(2) / 2, 0], + [-sqrt(2) / 2, sqrt(2) / 2, 0], [0, 0, 1]]) + assert A.ang_vel_in(N) == Vector(0) + + +def test_deprecated_joint_pos(): + N, A, P, C = _generate_body() + with warns_deprecated_sympy(): + pin = PinJoint('J', P, C, parent_joint_pos=N.x + N.y, + child_joint_pos=C.y - C.z) + assert pin.parent_point.pos_from(P.masscenter) == N.x + N.y + assert pin.child_point.pos_from(C.masscenter) == C.y - C.z + + N, A, P, C = _generate_body() + with warns_deprecated_sympy(): + slider = PrismaticJoint('J', P, C, parent_joint_pos=N.z + N.y, + child_joint_pos=C.y - C.x) + assert slider.parent_point.pos_from(P.masscenter) == N.z + N.y + assert slider.child_point.pos_from(C.masscenter) == C.y - C.x diff --git a/lib/python3.12/site-packages/sympy/physics/mechanics/tests/test_jointsmethod.py b/lib/python3.12/site-packages/sympy/physics/mechanics/tests/test_jointsmethod.py new file mode 100644 index 0000000000000000000000000000000000000000..1b48eae06dadc627442fd4e42445450be0393e33 --- /dev/null +++ b/lib/python3.12/site-packages/sympy/physics/mechanics/tests/test_jointsmethod.py @@ -0,0 +1,249 @@ +from sympy.core.function import expand +from sympy.core.symbol import symbols +from sympy.functions.elementary.trigonometric import (cos, sin) +from sympy.matrices.dense import Matrix +from sympy.simplify.trigsimp import trigsimp +from sympy.physics.mechanics import ( + PinJoint, JointsMethod, RigidBody, Particle, Body, KanesMethod, + PrismaticJoint, LagrangesMethod, inertia) +from sympy.physics.vector import dynamicsymbols, ReferenceFrame +from sympy.testing.pytest import raises, warns_deprecated_sympy +from sympy import zeros +from sympy.utilities.lambdify import lambdify +from sympy.solvers.solvers import solve + + +t = dynamicsymbols._t # type: ignore + + +def test_jointsmethod(): + with warns_deprecated_sympy(): + P = Body('P') + C = Body('C') + Pin = PinJoint('P1', P, C) + C_ixx, g = symbols('C_ixx g') + q, u = dynamicsymbols('q_P1, u_P1') + P.apply_force(g*P.y) + with warns_deprecated_sympy(): + method = JointsMethod(P, Pin) + assert method.frame == P.frame + assert method.bodies == [C, P] + assert method.loads == [(P.masscenter, g*P.frame.y)] + assert method.q == Matrix([q]) + assert method.u == Matrix([u]) + assert method.kdes == Matrix([u - q.diff()]) + soln = method.form_eoms() + assert soln == Matrix([[-C_ixx*u.diff()]]) + assert method.forcing_full == Matrix([[u], [0]]) + assert method.mass_matrix_full == Matrix([[1, 0], [0, C_ixx]]) + assert isinstance(method.method, KanesMethod) + + +def test_rigid_body_particle_compatibility(): + l, m, g = symbols('l m g') + C = RigidBody('C') + b = Particle('b', mass=m) + b_frame = ReferenceFrame('b_frame') + q, u = dynamicsymbols('q u') + P = PinJoint('P', C, b, coordinates=q, speeds=u, child_interframe=b_frame, + child_point=-l * b_frame.x, joint_axis=C.z) + with warns_deprecated_sympy(): + method = JointsMethod(C, P) + method.loads.append((b.masscenter, m * g * C.x)) + method.form_eoms() + rhs = method.rhs() + assert rhs[1] == -g*sin(q)/l + + +def test_jointmethod_duplicate_coordinates_speeds(): + with warns_deprecated_sympy(): + P = Body('P') + C = Body('C') + T = Body('T') + q, u = dynamicsymbols('q u') + P1 = PinJoint('P1', P, C, q) + P2 = PrismaticJoint('P2', C, T, q) + with warns_deprecated_sympy(): + raises(ValueError, lambda: JointsMethod(P, P1, P2)) + + P1 = PinJoint('P1', P, C, speeds=u) + P2 = PrismaticJoint('P2', C, T, speeds=u) + with warns_deprecated_sympy(): + raises(ValueError, lambda: JointsMethod(P, P1, P2)) + + P1 = PinJoint('P1', P, C, q, u) + P2 = PrismaticJoint('P2', C, T, q, u) + with warns_deprecated_sympy(): + raises(ValueError, lambda: JointsMethod(P, P1, P2)) + +def test_complete_simple_double_pendulum(): + q1, q2 = dynamicsymbols('q1 q2') + u1, u2 = dynamicsymbols('u1 u2') + m, l, g = symbols('m l g') + with warns_deprecated_sympy(): + C = Body('C') # ceiling + PartP = Body('P', mass=m) + PartR = Body('R', mass=m) + J1 = PinJoint('J1', C, PartP, speeds=u1, coordinates=q1, + child_point=-l*PartP.x, joint_axis=C.z) + J2 = PinJoint('J2', PartP, PartR, speeds=u2, coordinates=q2, + child_point=-l*PartR.x, joint_axis=PartP.z) + + PartP.apply_force(m*g*C.x) + PartR.apply_force(m*g*C.x) + + with warns_deprecated_sympy(): + method = JointsMethod(C, J1, J2) + method.form_eoms() + + assert expand(method.mass_matrix_full) == Matrix([[1, 0, 0, 0], + [0, 1, 0, 0], + [0, 0, 2*l**2*m*cos(q2) + 3*l**2*m, l**2*m*cos(q2) + l**2*m], + [0, 0, l**2*m*cos(q2) + l**2*m, l**2*m]]) + assert trigsimp(method.forcing_full) == trigsimp(Matrix([[u1], [u2], [-g*l*m*(sin(q1 + q2) + sin(q1)) - + g*l*m*sin(q1) + l**2*m*(2*u1 + u2)*u2*sin(q2)], + [-g*l*m*sin(q1 + q2) - l**2*m*u1**2*sin(q2)]])) + +def test_two_dof_joints(): + q1, q2, u1, u2 = dynamicsymbols('q1 q2 u1 u2') + m, c1, c2, k1, k2 = symbols('m c1 c2 k1 k2') + with warns_deprecated_sympy(): + W = Body('W') + B1 = Body('B1', mass=m) + B2 = Body('B2', mass=m) + J1 = PrismaticJoint('J1', W, B1, coordinates=q1, speeds=u1) + J2 = PrismaticJoint('J2', B1, B2, coordinates=q2, speeds=u2) + W.apply_force(k1*q1*W.x, reaction_body=B1) + W.apply_force(c1*u1*W.x, reaction_body=B1) + B1.apply_force(k2*q2*W.x, reaction_body=B2) + B1.apply_force(c2*u2*W.x, reaction_body=B2) + with warns_deprecated_sympy(): + method = JointsMethod(W, J1, J2) + method.form_eoms() + MM = method.mass_matrix + forcing = method.forcing + rhs = MM.LUsolve(forcing) + assert expand(rhs[0]) == expand((-k1 * q1 - c1 * u1 + k2 * q2 + c2 * u2)/m) + assert expand(rhs[1]) == expand((k1 * q1 + c1 * u1 - 2 * k2 * q2 - 2 * + c2 * u2) / m) + +def test_simple_pedulum(): + l, m, g = symbols('l m g') + with warns_deprecated_sympy(): + C = Body('C') + b = Body('b', mass=m) + q = dynamicsymbols('q') + P = PinJoint('P', C, b, speeds=q.diff(t), coordinates=q, + child_point=-l * b.x, joint_axis=C.z) + b.potential_energy = - m * g * l * cos(q) + with warns_deprecated_sympy(): + method = JointsMethod(C, P) + method.form_eoms(LagrangesMethod) + rhs = method.rhs() + assert rhs[1] == -g*sin(q)/l + +def test_chaos_pendulum(): + #https://www.pydy.org/examples/chaos_pendulum.html + mA, mB, lA, lB, IAxx, IBxx, IByy, IBzz, g = symbols('mA, mB, lA, lB, IAxx, IBxx, IByy, IBzz, g') + theta, phi, omega, alpha = dynamicsymbols('theta phi omega alpha') + + A = ReferenceFrame('A') + B = ReferenceFrame('B') + + with warns_deprecated_sympy(): + rod = Body('rod', mass=mA, frame=A, + central_inertia=inertia(A, IAxx, IAxx, 0)) + plate = Body('plate', mass=mB, frame=B, + central_inertia=inertia(B, IBxx, IByy, IBzz)) + C = Body('C') + J1 = PinJoint('J1', C, rod, coordinates=theta, speeds=omega, + child_point=-lA * rod.z, joint_axis=C.y) + J2 = PinJoint('J2', rod, plate, coordinates=phi, speeds=alpha, + parent_point=(lB - lA) * rod.z, joint_axis=rod.z) + + rod.apply_force(mA*g*C.z) + plate.apply_force(mB*g*C.z) + + with warns_deprecated_sympy(): + method = JointsMethod(C, J1, J2) + method.form_eoms() + + MM = method.mass_matrix + forcing = method.forcing + rhs = MM.LUsolve(forcing) + xd = (-2 * IBxx * alpha * omega * sin(phi) * cos(phi) + 2 * IByy * alpha * omega * sin(phi) * + cos(phi) - g * lA * mA * sin(theta) - g * lB * mB * sin(theta)) / (IAxx + IBxx * + sin(phi)**2 + IByy * cos(phi)**2 + lA**2 * mA + lB**2 * mB) + assert (rhs[0] - xd).simplify() == 0 + xd = (IBxx - IByy) * omega**2 * sin(phi) * cos(phi) / IBzz + assert (rhs[1] - xd).simplify() == 0 + +def test_four_bar_linkage_with_manual_constraints(): + q1, q2, q3, u1, u2, u3 = dynamicsymbols('q1:4, u1:4') + l1, l2, l3, l4, rho = symbols('l1:5, rho') + + N = ReferenceFrame('N') + inertias = [inertia(N, 0, 0, rho * l ** 3 / 12) for l in (l1, l2, l3, l4)] + with warns_deprecated_sympy(): + link1 = Body('Link1', frame=N, mass=rho * l1, + central_inertia=inertias[0]) + link2 = Body('Link2', mass=rho * l2, central_inertia=inertias[1]) + link3 = Body('Link3', mass=rho * l3, central_inertia=inertias[2]) + link4 = Body('Link4', mass=rho * l4, central_inertia=inertias[3]) + + joint1 = PinJoint( + 'J1', link1, link2, coordinates=q1, speeds=u1, joint_axis=link1.z, + parent_point=l1 / 2 * link1.x, child_point=-l2 / 2 * link2.x) + joint2 = PinJoint( + 'J2', link2, link3, coordinates=q2, speeds=u2, joint_axis=link2.z, + parent_point=l2 / 2 * link2.x, child_point=-l3 / 2 * link3.x) + joint3 = PinJoint( + 'J3', link3, link4, coordinates=q3, speeds=u3, joint_axis=link3.z, + parent_point=l3 / 2 * link3.x, child_point=-l4 / 2 * link4.x) + + loop = link4.masscenter.pos_from(link1.masscenter) \ + + l1 / 2 * link1.x + l4 / 2 * link4.x + + fh = Matrix([loop.dot(link1.x), loop.dot(link1.y)]) + + with warns_deprecated_sympy(): + method = JointsMethod(link1, joint1, joint2, joint3) + + t = dynamicsymbols._t + qdots = solve(method.kdes, [q1.diff(t), q2.diff(t), q3.diff(t)]) + fhd = fh.diff(t).subs(qdots) + + kane = KanesMethod(method.frame, q_ind=[q1], u_ind=[u1], + q_dependent=[q2, q3], u_dependent=[u2, u3], + kd_eqs=method.kdes, configuration_constraints=fh, + velocity_constraints=fhd, forcelist=method.loads, + bodies=method.bodies) + fr, frs = kane.kanes_equations() + assert fr == zeros(1) + + # Numerically check the mass- and forcing-matrix + p = Matrix([l1, l2, l3, l4, rho]) + q = Matrix([q1, q2, q3]) + u = Matrix([u1, u2, u3]) + eval_m = lambdify((q, p), kane.mass_matrix) + eval_f = lambdify((q, u, p), kane.forcing) + eval_fhd = lambdify((q, u, p), fhd) + + p_vals = [0.13, 0.24, 0.21, 0.34, 997] + q_vals = [2.1, 0.6655470375077588, 2.527408138024188] # Satisfies fh + u_vals = [0.2, -0.17963733938852067, 0.1309060540601612] # Satisfies fhd + mass_check = Matrix([[3.452709815256506e+01, 7.003948798374735e+00, + -4.939690970641498e+00], + [-2.203792703880936e-14, 2.071702479957077e-01, + 2.842917573033711e-01], + [-1.300000000000123e-01, -8.836934896046506e-03, + 1.864891330060847e-01]]) + forcing_check = Matrix([[-0.031211821321648], + [-0.00066022608181], + [0.001813559741243]]) + eps = 1e-10 + assert all(abs(x) < eps for x in eval_fhd(q_vals, u_vals, p_vals)) + assert all(abs(x) < eps for x in + (Matrix(eval_m(q_vals, p_vals)) - mass_check)) + assert all(abs(x) < eps for x in + (Matrix(eval_f(q_vals, u_vals, p_vals)) - forcing_check)) diff --git a/lib/python3.12/site-packages/sympy/physics/mechanics/tests/test_kane.py b/lib/python3.12/site-packages/sympy/physics/mechanics/tests/test_kane.py new file mode 100644 index 0000000000000000000000000000000000000000..5f9310aae6d720c32615a86df5e488f46a513c76 --- /dev/null +++ b/lib/python3.12/site-packages/sympy/physics/mechanics/tests/test_kane.py @@ -0,0 +1,553 @@ +from sympy import solve +from sympy import (cos, expand, Matrix, sin, symbols, tan, sqrt, S, + zeros, eye) +from sympy.simplify.simplify import simplify +from sympy.physics.mechanics import (dynamicsymbols, ReferenceFrame, Point, + RigidBody, KanesMethod, inertia, Particle, + dot, find_dynamicsymbols) +from sympy.testing.pytest import raises + + +def test_invalid_coordinates(): + # Simple pendulum, but use symbols instead of dynamicsymbols + l, m, g = symbols('l m g') + q, u = symbols('q u') # Generalized coordinate + kd = [q.diff(dynamicsymbols._t) - u] + N, O = ReferenceFrame('N'), Point('O') + O.set_vel(N, 0) + P = Particle('P', Point('P'), m) + P.point.set_pos(O, l * (sin(q) * N.x - cos(q) * N.y)) + F = (P.point, -m * g * N.y) + raises(ValueError, lambda: KanesMethod(N, [q], [u], kd, bodies=[P], + forcelist=[F])) + + +def test_one_dof(): + # This is for a 1 dof spring-mass-damper case. + # It is described in more detail in the KanesMethod docstring. + q, u = dynamicsymbols('q u') + qd, ud = dynamicsymbols('q u', 1) + m, c, k = symbols('m c k') + N = ReferenceFrame('N') + P = Point('P') + P.set_vel(N, u * N.x) + + kd = [qd - u] + FL = [(P, (-k * q - c * u) * N.x)] + pa = Particle('pa', P, m) + BL = [pa] + + KM = KanesMethod(N, [q], [u], kd) + KM.kanes_equations(BL, FL) + + assert KM.bodies == BL + assert KM.loads == FL + + MM = KM.mass_matrix + forcing = KM.forcing + rhs = MM.inv() * forcing + assert expand(rhs[0]) == expand(-(q * k + u * c) / m) + + assert simplify(KM.rhs() - + KM.mass_matrix_full.LUsolve(KM.forcing_full)) == zeros(2, 1) + + assert (KM.linearize(A_and_B=True, )[0] == Matrix([[0, 1], [-k/m, -c/m]])) + + +def test_two_dof(): + # This is for a 2 d.o.f., 2 particle spring-mass-damper. + # The first coordinate is the displacement of the first particle, and the + # second is the relative displacement between the first and second + # particles. Speeds are defined as the time derivatives of the particles. + q1, q2, u1, u2 = dynamicsymbols('q1 q2 u1 u2') + q1d, q2d, u1d, u2d = dynamicsymbols('q1 q2 u1 u2', 1) + m, c1, c2, k1, k2 = symbols('m c1 c2 k1 k2') + N = ReferenceFrame('N') + P1 = Point('P1') + P2 = Point('P2') + P1.set_vel(N, u1 * N.x) + P2.set_vel(N, (u1 + u2) * N.x) + # Note we multiply the kinematic equation by an arbitrary factor + # to test the implicit vs explicit kinematics attribute + kd = [q1d/2 - u1/2, 2*q2d - 2*u2] + + # Now we create the list of forces, then assign properties to each + # particle, then create a list of all particles. + FL = [(P1, (-k1 * q1 - c1 * u1 + k2 * q2 + c2 * u2) * N.x), (P2, (-k2 * + q2 - c2 * u2) * N.x)] + pa1 = Particle('pa1', P1, m) + pa2 = Particle('pa2', P2, m) + BL = [pa1, pa2] + + # Finally we create the KanesMethod object, specify the inertial frame, + # pass relevant information, and form Fr & Fr*. Then we calculate the mass + # matrix and forcing terms, and finally solve for the udots. + KM = KanesMethod(N, q_ind=[q1, q2], u_ind=[u1, u2], kd_eqs=kd) + KM.kanes_equations(BL, FL) + MM = KM.mass_matrix + forcing = KM.forcing + rhs = MM.inv() * forcing + assert expand(rhs[0]) == expand((-k1 * q1 - c1 * u1 + k2 * q2 + c2 * u2)/m) + assert expand(rhs[1]) == expand((k1 * q1 + c1 * u1 - 2 * k2 * q2 - 2 * + c2 * u2) / m) + + # Check that the explicit form is the default and kinematic mass matrix is identity + assert KM.explicit_kinematics + assert KM.mass_matrix_kin == eye(2) + + # Check that for the implicit form the mass matrix is not identity + KM.explicit_kinematics = False + assert KM.mass_matrix_kin == Matrix([[S(1)/2, 0], [0, 2]]) + + # Check that whether using implicit or explicit kinematics the RHS + # equations are consistent with the matrix form + for explicit_kinematics in [False, True]: + KM.explicit_kinematics = explicit_kinematics + assert simplify(KM.rhs() - + KM.mass_matrix_full.LUsolve(KM.forcing_full)) == zeros(4, 1) + + # Make sure an error is raised if nonlinear kinematic differential + # equations are supplied. + kd = [q1d - u1**2, sin(q2d) - cos(u2)] + raises(ValueError, lambda: KanesMethod(N, q_ind=[q1, q2], + u_ind=[u1, u2], kd_eqs=kd)) + +def test_pend(): + q, u = dynamicsymbols('q u') + qd, ud = dynamicsymbols('q u', 1) + m, l, g = symbols('m l g') + N = ReferenceFrame('N') + P = Point('P') + P.set_vel(N, -l * u * sin(q) * N.x + l * u * cos(q) * N.y) + kd = [qd - u] + + FL = [(P, m * g * N.x)] + pa = Particle('pa', P, m) + BL = [pa] + + KM = KanesMethod(N, [q], [u], kd) + KM.kanes_equations(BL, FL) + MM = KM.mass_matrix + forcing = KM.forcing + rhs = MM.inv() * forcing + rhs.simplify() + assert expand(rhs[0]) == expand(-g / l * sin(q)) + assert simplify(KM.rhs() - + KM.mass_matrix_full.LUsolve(KM.forcing_full)) == zeros(2, 1) + + +def test_rolling_disc(): + # Rolling Disc Example + # Here the rolling disc is formed from the contact point up, removing the + # need to introduce generalized speeds. Only 3 configuration and three + # speed variables are need to describe this system, along with the disc's + # mass and radius, and the local gravity (note that mass will drop out). + q1, q2, q3, u1, u2, u3 = dynamicsymbols('q1 q2 q3 u1 u2 u3') + q1d, q2d, q3d, u1d, u2d, u3d = dynamicsymbols('q1 q2 q3 u1 u2 u3', 1) + r, m, g = symbols('r m g') + + # The kinematics are formed by a series of simple rotations. Each simple + # rotation creates a new frame, and the next rotation is defined by the new + # frame's basis vectors. This example uses a 3-1-2 series of rotations, or + # Z, X, Y series of rotations. Angular velocity for this is defined using + # the second frame's basis (the lean frame). + N = ReferenceFrame('N') + Y = N.orientnew('Y', 'Axis', [q1, N.z]) + L = Y.orientnew('L', 'Axis', [q2, Y.x]) + R = L.orientnew('R', 'Axis', [q3, L.y]) + w_R_N_qd = R.ang_vel_in(N) + R.set_ang_vel(N, u1 * L.x + u2 * L.y + u3 * L.z) + + # This is the translational kinematics. We create a point with no velocity + # in N; this is the contact point between the disc and ground. Next we form + # the position vector from the contact point to the disc's center of mass. + # Finally we form the velocity and acceleration of the disc. + C = Point('C') + C.set_vel(N, 0) + Dmc = C.locatenew('Dmc', r * L.z) + Dmc.v2pt_theory(C, N, R) + + # This is a simple way to form the inertia dyadic. + I = inertia(L, m / 4 * r**2, m / 2 * r**2, m / 4 * r**2) + + # Kinematic differential equations; how the generalized coordinate time + # derivatives relate to generalized speeds. + kd = [dot(R.ang_vel_in(N) - w_R_N_qd, uv) for uv in L] + + # Creation of the force list; it is the gravitational force at the mass + # center of the disc. Then we create the disc by assigning a Point to the + # center of mass attribute, a ReferenceFrame to the frame attribute, and mass + # and inertia. Then we form the body list. + ForceList = [(Dmc, - m * g * Y.z)] + BodyD = RigidBody('BodyD', Dmc, R, m, (I, Dmc)) + BodyList = [BodyD] + + # Finally we form the equations of motion, using the same steps we did + # before. Specify inertial frame, supply generalized speeds, supply + # kinematic differential equation dictionary, compute Fr from the force + # list and Fr* from the body list, compute the mass matrix and forcing + # terms, then solve for the u dots (time derivatives of the generalized + # speeds). + KM = KanesMethod(N, q_ind=[q1, q2, q3], u_ind=[u1, u2, u3], kd_eqs=kd) + KM.kanes_equations(BodyList, ForceList) + MM = KM.mass_matrix + forcing = KM.forcing + rhs = MM.inv() * forcing + kdd = KM.kindiffdict() + rhs = rhs.subs(kdd) + rhs.simplify() + assert rhs.expand() == Matrix([(6*u2*u3*r - u3**2*r*tan(q2) + + 4*g*sin(q2))/(5*r), -2*u1*u3/3, u1*(-2*u2 + u3*tan(q2))]).expand() + assert simplify(KM.rhs() - + KM.mass_matrix_full.LUsolve(KM.forcing_full)) == zeros(6, 1) + + # This code tests our output vs. benchmark values. When r=g=m=1, the + # critical speed (where all eigenvalues of the linearized equations are 0) + # is 1 / sqrt(3) for the upright case. + A = KM.linearize(A_and_B=True)[0] + A_upright = A.subs({r: 1, g: 1, m: 1}).subs({q1: 0, q2: 0, q3: 0, u1: 0, u3: 0}) + import sympy + assert sympy.sympify(A_upright.subs({u2: 1 / sqrt(3)})).eigenvals() == {S.Zero: 6} + + +def test_aux(): + # Same as above, except we have 2 auxiliary speeds for the ground contact + # point, which is known to be zero. In one case, we go through then + # substitute the aux. speeds in at the end (they are zero, as well as their + # derivative), in the other case, we use the built-in auxiliary speed part + # of KanesMethod. The equations from each should be the same. + q1, q2, q3, u1, u2, u3 = dynamicsymbols('q1 q2 q3 u1 u2 u3') + q1d, q2d, q3d, u1d, u2d, u3d = dynamicsymbols('q1 q2 q3 u1 u2 u3', 1) + u4, u5, f1, f2 = dynamicsymbols('u4, u5, f1, f2') + u4d, u5d = dynamicsymbols('u4, u5', 1) + r, m, g = symbols('r m g') + + N = ReferenceFrame('N') + Y = N.orientnew('Y', 'Axis', [q1, N.z]) + L = Y.orientnew('L', 'Axis', [q2, Y.x]) + R = L.orientnew('R', 'Axis', [q3, L.y]) + w_R_N_qd = R.ang_vel_in(N) + R.set_ang_vel(N, u1 * L.x + u2 * L.y + u3 * L.z) + + C = Point('C') + C.set_vel(N, u4 * L.x + u5 * (Y.z ^ L.x)) + Dmc = C.locatenew('Dmc', r * L.z) + Dmc.v2pt_theory(C, N, R) + Dmc.a2pt_theory(C, N, R) + + I = inertia(L, m / 4 * r**2, m / 2 * r**2, m / 4 * r**2) + + kd = [dot(R.ang_vel_in(N) - w_R_N_qd, uv) for uv in L] + + ForceList = [(Dmc, - m * g * Y.z), (C, f1 * L.x + f2 * (Y.z ^ L.x))] + BodyD = RigidBody('BodyD', Dmc, R, m, (I, Dmc)) + BodyList = [BodyD] + + KM = KanesMethod(N, q_ind=[q1, q2, q3], u_ind=[u1, u2, u3, u4, u5], + kd_eqs=kd) + (fr, frstar) = KM.kanes_equations(BodyList, ForceList) + fr = fr.subs({u4d: 0, u5d: 0}).subs({u4: 0, u5: 0}) + frstar = frstar.subs({u4d: 0, u5d: 0}).subs({u4: 0, u5: 0}) + + KM2 = KanesMethod(N, q_ind=[q1, q2, q3], u_ind=[u1, u2, u3], kd_eqs=kd, + u_auxiliary=[u4, u5]) + (fr2, frstar2) = KM2.kanes_equations(BodyList, ForceList) + fr2 = fr2.subs({u4d: 0, u5d: 0}).subs({u4: 0, u5: 0}) + frstar2 = frstar2.subs({u4d: 0, u5d: 0}).subs({u4: 0, u5: 0}) + + frstar.simplify() + frstar2.simplify() + + assert (fr - fr2).expand() == Matrix([0, 0, 0, 0, 0]) + assert (frstar - frstar2).expand() == Matrix([0, 0, 0, 0, 0]) + + +def test_parallel_axis(): + # This is for a 2 dof inverted pendulum on a cart. + # This tests the parallel axis code in KanesMethod. The inertia of the + # pendulum is defined about the hinge, not about the center of mass. + + # Defining the constants and knowns of the system + gravity = symbols('g') + k, ls = symbols('k ls') + a, mA, mC = symbols('a mA mC') + F = dynamicsymbols('F') + Ix, Iy, Iz = symbols('Ix Iy Iz') + + # Declaring the Generalized coordinates and speeds + q1, q2 = dynamicsymbols('q1 q2') + q1d, q2d = dynamicsymbols('q1 q2', 1) + u1, u2 = dynamicsymbols('u1 u2') + u1d, u2d = dynamicsymbols('u1 u2', 1) + + # Creating reference frames + N = ReferenceFrame('N') + A = ReferenceFrame('A') + + A.orient(N, 'Axis', [-q2, N.z]) + A.set_ang_vel(N, -u2 * N.z) + + # Origin of Newtonian reference frame + O = Point('O') + + # Creating and Locating the positions of the cart, C, and the + # center of mass of the pendulum, A + C = O.locatenew('C', q1 * N.x) + Ao = C.locatenew('Ao', a * A.y) + + # Defining velocities of the points + O.set_vel(N, 0) + C.set_vel(N, u1 * N.x) + Ao.v2pt_theory(C, N, A) + Cart = Particle('Cart', C, mC) + Pendulum = RigidBody('Pendulum', Ao, A, mA, (inertia(A, Ix, Iy, Iz), C)) + + # kinematical differential equations + + kindiffs = [q1d - u1, q2d - u2] + + bodyList = [Cart, Pendulum] + + forceList = [(Ao, -N.y * gravity * mA), + (C, -N.y * gravity * mC), + (C, -N.x * k * (q1 - ls)), + (C, N.x * F)] + + km = KanesMethod(N, [q1, q2], [u1, u2], kindiffs) + (fr, frstar) = km.kanes_equations(bodyList, forceList) + mm = km.mass_matrix_full + assert mm[3, 3] == Iz + +def test_input_format(): + # 1 dof problem from test_one_dof + q, u = dynamicsymbols('q u') + qd, ud = dynamicsymbols('q u', 1) + m, c, k = symbols('m c k') + N = ReferenceFrame('N') + P = Point('P') + P.set_vel(N, u * N.x) + + kd = [qd - u] + FL = [(P, (-k * q - c * u) * N.x)] + pa = Particle('pa', P, m) + BL = [pa] + + KM = KanesMethod(N, [q], [u], kd) + # test for input format kane.kanes_equations((body1, body2, particle1)) + assert KM.kanes_equations(BL)[0] == Matrix([0]) + # test for input format kane.kanes_equations(bodies=(body1, body 2), loads=(load1,load2)) + assert KM.kanes_equations(bodies=BL, loads=None)[0] == Matrix([0]) + # test for input format kane.kanes_equations(bodies=(body1, body 2), loads=None) + assert KM.kanes_equations(BL, loads=None)[0] == Matrix([0]) + # test for input format kane.kanes_equations(bodies=(body1, body 2)) + assert KM.kanes_equations(BL)[0] == Matrix([0]) + # test for input format kane.kanes_equations(bodies=(body1, body2), loads=[]) + assert KM.kanes_equations(BL, [])[0] == Matrix([0]) + # test for error raised when a wrong force list (in this case a string) is provided + raises(ValueError, lambda: KM._form_fr('bad input')) + + # 1 dof problem from test_one_dof with FL & BL in instance + KM = KanesMethod(N, [q], [u], kd, bodies=BL, forcelist=FL) + assert KM.kanes_equations()[0] == Matrix([-c*u - k*q]) + + # 2 dof problem from test_two_dof + q1, q2, u1, u2 = dynamicsymbols('q1 q2 u1 u2') + q1d, q2d, u1d, u2d = dynamicsymbols('q1 q2 u1 u2', 1) + m, c1, c2, k1, k2 = symbols('m c1 c2 k1 k2') + N = ReferenceFrame('N') + P1 = Point('P1') + P2 = Point('P2') + P1.set_vel(N, u1 * N.x) + P2.set_vel(N, (u1 + u2) * N.x) + kd = [q1d - u1, q2d - u2] + + FL = ((P1, (-k1 * q1 - c1 * u1 + k2 * q2 + c2 * u2) * N.x), (P2, (-k2 * + q2 - c2 * u2) * N.x)) + pa1 = Particle('pa1', P1, m) + pa2 = Particle('pa2', P2, m) + BL = (pa1, pa2) + + KM = KanesMethod(N, q_ind=[q1, q2], u_ind=[u1, u2], kd_eqs=kd) + # test for input format + # kane.kanes_equations((body1, body2), (load1, load2)) + KM.kanes_equations(BL, FL) + MM = KM.mass_matrix + forcing = KM.forcing + rhs = MM.inv() * forcing + assert expand(rhs[0]) == expand((-k1 * q1 - c1 * u1 + k2 * q2 + c2 * u2)/m) + assert expand(rhs[1]) == expand((k1 * q1 + c1 * u1 - 2 * k2 * q2 - 2 * + c2 * u2) / m) + + +def test_implicit_kinematics(): + # Test that implicit kinematics can handle complicated + # equations that explicit form struggles with + # See https://github.com/sympy/sympy/issues/22626 + + # Inertial frame + NED = ReferenceFrame('NED') + NED_o = Point('NED_o') + NED_o.set_vel(NED, 0) + + # body frame + q_att = dynamicsymbols('lambda_0:4', real=True) + B = NED.orientnew('B', 'Quaternion', q_att) + + # Generalized coordinates + q_pos = dynamicsymbols('B_x:z') + B_cm = NED_o.locatenew('B_cm', q_pos[0]*B.x + q_pos[1]*B.y + q_pos[2]*B.z) + + q_ind = q_att[1:] + q_pos + q_dep = [q_att[0]] + + kinematic_eqs = [] + + # Generalized velocities + B_ang_vel = B.ang_vel_in(NED) + P, Q, R = dynamicsymbols('P Q R') + B.set_ang_vel(NED, P*B.x + Q*B.y + R*B.z) + + B_ang_vel_kd = (B.ang_vel_in(NED) - B_ang_vel).simplify() + + # Equating the two gives us the kinematic equation + kinematic_eqs += [ + B_ang_vel_kd & B.x, + B_ang_vel_kd & B.y, + B_ang_vel_kd & B.z + ] + + B_cm_vel = B_cm.vel(NED) + U, V, W = dynamicsymbols('U V W') + B_cm.set_vel(NED, U*B.x + V*B.y + W*B.z) + + # Compute the velocity of the point using the two methods + B_ref_vel_kd = (B_cm.vel(NED) - B_cm_vel) + + # taking dot product with unit vectors to get kinematic equations + # relating body coordinates and velocities + + # Note, there is a choice to dot with NED.xyz here. That makes + # the implicit form have some bigger terms but is still fine, the + # explicit form still struggles though + kinematic_eqs += [ + B_ref_vel_kd & B.x, + B_ref_vel_kd & B.y, + B_ref_vel_kd & B.z, + ] + + u_ind = [U, V, W, P, Q, R] + + # constraints + q_att_vec = Matrix(q_att) + config_cons = [(q_att_vec.T*q_att_vec)[0] - 1] #unit norm + kinematic_eqs = kinematic_eqs + [(q_att_vec.T * q_att_vec.diff())[0]] + + try: + KM = KanesMethod(NED, q_ind, u_ind, + q_dependent= q_dep, + kd_eqs = kinematic_eqs, + configuration_constraints = config_cons, + velocity_constraints= [], + u_dependent= [], #no dependent speeds + u_auxiliary = [], # No auxiliary speeds + explicit_kinematics = False # implicit kinematics + ) + except Exception as e: + raise e + + # mass and inertia dyadic relative to CM + M_B = symbols('M_B') + J_B = inertia(B, *[S(f'J_B_{ax}')*(1 if ax[0] == ax[1] else -1) + for ax in ['xx', 'yy', 'zz', 'xy', 'yz', 'xz']]) + J_B = J_B.subs({S('J_B_xy'): 0, S('J_B_yz'): 0}) + RB = RigidBody('RB', B_cm, B, M_B, (J_B, B_cm)) + + rigid_bodies = [RB] + # Forces + force_list = [ + #gravity pointing down + (RB.masscenter, RB.mass*S('g')*NED.z), + #generic forces and torques in body frame(inputs) + (RB.frame, dynamicsymbols('T_z')*B.z), + (RB.masscenter, dynamicsymbols('F_z')*B.z) + ] + + KM.kanes_equations(rigid_bodies, force_list) + + # Expecting implicit form to be less than 5% of the flops + n_ops_implicit = sum( + [x.count_ops() for x in KM.forcing_full] + + [x.count_ops() for x in KM.mass_matrix_full] + ) + # Save implicit kinematic matrices to use later + mass_matrix_kin_implicit = KM.mass_matrix_kin + forcing_kin_implicit = KM.forcing_kin + + KM.explicit_kinematics = True + n_ops_explicit = sum( + [x.count_ops() for x in KM.forcing_full] + + [x.count_ops() for x in KM.mass_matrix_full] + ) + forcing_kin_explicit = KM.forcing_kin + + assert n_ops_implicit / n_ops_explicit < .05 + + # Ideally we would check that implicit and explicit equations give the same result as done in test_one_dof + # But the whole raison-d'etre of the implicit equations is to deal with problems such + # as this one where the explicit form is too complicated to handle, especially the angular part + # (i.e. tests would be too slow) + # Instead, we check that the kinematic equations are correct using more fundamental tests: + # + # (1) that we recover the kinematic equations we have provided + assert (mass_matrix_kin_implicit * KM.q.diff() - forcing_kin_implicit) == Matrix(kinematic_eqs) + + # (2) that rate of quaternions matches what 'textbook' solutions give + # Note that we just use the explicit kinematics for the linear velocities + # as they are not as complicated as the angular ones + qdot_candidate = forcing_kin_explicit + + quat_dot_textbook = Matrix([ + [0, -P, -Q, -R], + [P, 0, R, -Q], + [Q, -R, 0, P], + [R, Q, -P, 0], + ]) * q_att_vec / 2 + + # Again, if we don't use this "textbook" solution + # sympy will struggle to deal with the terms related to quaternion rates + # due to the number of operations involved + qdot_candidate[-1] = quat_dot_textbook[0] # lambda_0, note the [-1] as sympy's Kane puts the dependent coordinate last + qdot_candidate[0] = quat_dot_textbook[1] # lambda_1 + qdot_candidate[1] = quat_dot_textbook[2] # lambda_2 + qdot_candidate[2] = quat_dot_textbook[3] # lambda_3 + + # sub the config constraint in the candidate solution and compare to the implicit rhs + lambda_0_sol = solve(config_cons[0], q_att_vec[0])[1] + lhs_candidate = simplify(mass_matrix_kin_implicit * qdot_candidate).subs({q_att_vec[0]: lambda_0_sol}) + assert lhs_candidate == forcing_kin_implicit + +def test_issue_24887(): + # Spherical pendulum + g, l, m, c = symbols('g l m c') + q1, q2, q3, u1, u2, u3 = dynamicsymbols('q1:4 u1:4') + N = ReferenceFrame('N') + A = ReferenceFrame('A') + A.orient_body_fixed(N, (q1, q2, q3), 'zxy') + N_w_A = A.ang_vel_in(N) + # A.set_ang_vel(N, u1 * A.x + u2 * A.y + u3 * A.z) + kdes = [N_w_A.dot(A.x) - u1, N_w_A.dot(A.y) - u2, N_w_A.dot(A.z) - u3] + O = Point('O') + O.set_vel(N, 0) + Po = O.locatenew('Po', -l * A.y) + Po.set_vel(A, 0) + P = Particle('P', Po, m) + kane = KanesMethod(N, [q1, q2, q3], [u1, u2, u3], kdes, bodies=[P], + forcelist=[(Po, -m * g * N.y)]) + kane.kanes_equations() + expected_md = m * l ** 2 * Matrix([[1, 0, 0], [0, 0, 0], [0, 0, 1]]) + expected_fd = Matrix([ + [l*m*(g*(sin(q1)*sin(q3) - sin(q2)*cos(q1)*cos(q3)) - l*u2*u3)], + [0], [l*m*(-g*(sin(q1)*cos(q3) + sin(q2)*sin(q3)*cos(q1)) + l*u1*u2)]]) + assert find_dynamicsymbols(kane.forcing).issubset({q1, q2, q3, u1, u2, u3}) + assert simplify(kane.mass_matrix - expected_md) == zeros(3, 3) + assert simplify(kane.forcing - expected_fd) == zeros(3, 1) diff --git a/lib/python3.12/site-packages/sympy/physics/mechanics/tests/test_kane2.py b/lib/python3.12/site-packages/sympy/physics/mechanics/tests/test_kane2.py new file mode 100644 index 0000000000000000000000000000000000000000..e55866672aec0adcd951e772964a4ed205b56405 --- /dev/null +++ b/lib/python3.12/site-packages/sympy/physics/mechanics/tests/test_kane2.py @@ -0,0 +1,464 @@ +from sympy import cos, Matrix, sin, zeros, tan, pi, symbols +from sympy.simplify.simplify import simplify +from sympy.simplify.trigsimp import trigsimp +from sympy.solvers.solvers import solve +from sympy.physics.mechanics import (cross, dot, dynamicsymbols, + find_dynamicsymbols, KanesMethod, inertia, + inertia_of_point_mass, Point, + ReferenceFrame, RigidBody) + + +def test_aux_dep(): + # This test is about rolling disc dynamics, comparing the results found + # with KanesMethod to those found when deriving the equations "manually" + # with SymPy. + # The terms Fr, Fr*, and Fr*_steady are all compared between the two + # methods. Here, Fr*_steady refers to the generalized inertia forces for an + # equilibrium configuration. + # Note: comparing to the test of test_rolling_disc() in test_kane.py, this + # test also tests auxiliary speeds and configuration and motion constraints + #, seen in the generalized dependent coordinates q[3], and depend speeds + # u[3], u[4] and u[5]. + + + # First, manual derivation of Fr, Fr_star, Fr_star_steady. + + # Symbols for time and constant parameters. + # Symbols for contact forces: Fx, Fy, Fz. + t, r, m, g, I, J = symbols('t r m g I J') + Fx, Fy, Fz = symbols('Fx Fy Fz') + + # Configuration variables and their time derivatives: + # q[0] -- yaw + # q[1] -- lean + # q[2] -- spin + # q[3] -- dot(-r*B.z, A.z) -- distance from ground plane to disc center in + # A.z direction + # Generalized speeds and their time derivatives: + # u[0] -- disc angular velocity component, disc fixed x direction + # u[1] -- disc angular velocity component, disc fixed y direction + # u[2] -- disc angular velocity component, disc fixed z direction + # u[3] -- disc velocity component, A.x direction + # u[4] -- disc velocity component, A.y direction + # u[5] -- disc velocity component, A.z direction + # Auxiliary generalized speeds: + # ua[0] -- contact point auxiliary generalized speed, A.x direction + # ua[1] -- contact point auxiliary generalized speed, A.y direction + # ua[2] -- contact point auxiliary generalized speed, A.z direction + q = dynamicsymbols('q:4') + qd = [qi.diff(t) for qi in q] + u = dynamicsymbols('u:6') + ud = [ui.diff(t) for ui in u] + ud_zero = dict(zip(ud, [0.]*len(ud))) + ua = dynamicsymbols('ua:3') + ua_zero = dict(zip(ua, [0.]*len(ua))) # noqa:F841 + + # Reference frames: + # Yaw intermediate frame: A. + # Lean intermediate frame: B. + # Disc fixed frame: C. + N = ReferenceFrame('N') + A = N.orientnew('A', 'Axis', [q[0], N.z]) + B = A.orientnew('B', 'Axis', [q[1], A.x]) + C = B.orientnew('C', 'Axis', [q[2], B.y]) + + # Angular velocity and angular acceleration of disc fixed frame + # u[0], u[1] and u[2] are generalized independent speeds. + C.set_ang_vel(N, u[0]*B.x + u[1]*B.y + u[2]*B.z) + C.set_ang_acc(N, C.ang_vel_in(N).diff(t, B) + + cross(B.ang_vel_in(N), C.ang_vel_in(N))) + + # Velocity and acceleration of points: + # Disc-ground contact point: P. + # Center of disc: O, defined from point P with depend coordinate: q[3] + # u[3], u[4] and u[5] are generalized dependent speeds. + P = Point('P') + P.set_vel(N, ua[0]*A.x + ua[1]*A.y + ua[2]*A.z) + O = P.locatenew('O', q[3]*A.z + r*sin(q[1])*A.y) + O.set_vel(N, u[3]*A.x + u[4]*A.y + u[5]*A.z) + O.set_acc(N, O.vel(N).diff(t, A) + cross(A.ang_vel_in(N), O.vel(N))) + + # Kinematic differential equations: + # Two equalities: one is w_c_n_qd = C.ang_vel_in(N) in three coordinates + # directions of B, for qd0, qd1 and qd2. + # the other is v_o_n_qd = O.vel(N) in A.z direction for qd3. + # Then, solve for dq/dt's in terms of u's: qd_kd. + w_c_n_qd = qd[0]*A.z + qd[1]*B.x + qd[2]*B.y + v_o_n_qd = O.pos_from(P).diff(t, A) + cross(A.ang_vel_in(N), O.pos_from(P)) + kindiffs = Matrix([dot(w_c_n_qd - C.ang_vel_in(N), uv) for uv in B] + + [dot(v_o_n_qd - O.vel(N), A.z)]) + qd_kd = solve(kindiffs, qd) # noqa:F841 + + # Values of generalized speeds during a steady turn for later substitution + # into the Fr_star_steady. + steady_conditions = solve(kindiffs.subs({qd[1] : 0, qd[3] : 0}), u) + steady_conditions.update({qd[1] : 0, qd[3] : 0}) + + # Partial angular velocities and velocities. + partial_w_C = [C.ang_vel_in(N).diff(ui, N) for ui in u + ua] + partial_v_O = [O.vel(N).diff(ui, N) for ui in u + ua] + partial_v_P = [P.vel(N).diff(ui, N) for ui in u + ua] + + # Configuration constraint: f_c, the projection of radius r in A.z direction + # is q[3]. + # Velocity constraints: f_v, for u3, u4 and u5. + # Acceleration constraints: f_a. + f_c = Matrix([dot(-r*B.z, A.z) - q[3]]) + f_v = Matrix([dot(O.vel(N) - (P.vel(N) + cross(C.ang_vel_in(N), + O.pos_from(P))), ai).expand() for ai in A]) + v_o_n = cross(C.ang_vel_in(N), O.pos_from(P)) + a_o_n = v_o_n.diff(t, A) + cross(A.ang_vel_in(N), v_o_n) + f_a = Matrix([dot(O.acc(N) - a_o_n, ai) for ai in A]) # noqa:F841 + + # Solve for constraint equations in the form of + # u_dependent = A_rs * [u_i; u_aux]. + # First, obtain constraint coefficient matrix: M_v * [u; ua] = 0; + # Second, taking u[0], u[1], u[2] as independent, + # taking u[3], u[4], u[5] as dependent, + # rearranging the matrix of M_v to be A_rs for u_dependent. + # Third, u_aux ==0 for u_dep, and resulting dictionary of u_dep_dict. + M_v = zeros(3, 9) + for i in range(3): + for j, ui in enumerate(u + ua): + M_v[i, j] = f_v[i].diff(ui) + + M_v_i = M_v[:, :3] + M_v_d = M_v[:, 3:6] + M_v_aux = M_v[:, 6:] + M_v_i_aux = M_v_i.row_join(M_v_aux) + A_rs = - M_v_d.inv() * M_v_i_aux + + u_dep = A_rs[:, :3] * Matrix(u[:3]) + u_dep_dict = dict(zip(u[3:], u_dep)) + + # Active forces: F_O acting on point O; F_P acting on point P. + # Generalized active forces (unconstrained): Fr_u = F_point * pv_point. + F_O = m*g*A.z + F_P = Fx * A.x + Fy * A.y + Fz * A.z + Fr_u = Matrix([dot(F_O, pv_o) + dot(F_P, pv_p) for pv_o, pv_p in + zip(partial_v_O, partial_v_P)]) + + # Inertia force: R_star_O. + # Inertia of disc: I_C_O, where J is a inertia component about principal axis. + # Inertia torque: T_star_C. + # Generalized inertia forces (unconstrained): Fr_star_u. + R_star_O = -m*O.acc(N) + I_C_O = inertia(B, I, J, I) + T_star_C = -(dot(I_C_O, C.ang_acc_in(N)) \ + + cross(C.ang_vel_in(N), dot(I_C_O, C.ang_vel_in(N)))) + Fr_star_u = Matrix([dot(R_star_O, pv) + dot(T_star_C, pav) for pv, pav in + zip(partial_v_O, partial_w_C)]) + + # Form nonholonomic Fr: Fr_c, and nonholonomic Fr_star: Fr_star_c. + # Also, nonholonomic Fr_star in steady turning condition: Fr_star_steady. + Fr_c = Fr_u[:3, :].col_join(Fr_u[6:, :]) + A_rs.T * Fr_u[3:6, :] + Fr_star_c = Fr_star_u[:3, :].col_join(Fr_star_u[6:, :])\ + + A_rs.T * Fr_star_u[3:6, :] + Fr_star_steady = Fr_star_c.subs(ud_zero).subs(u_dep_dict)\ + .subs(steady_conditions).subs({q[3]: -r*cos(q[1])}).expand() + + + # Second, using KaneMethod in mechanics for fr, frstar and frstar_steady. + + # Rigid Bodies: disc, with inertia I_C_O. + iner_tuple = (I_C_O, O) + disc = RigidBody('disc', O, C, m, iner_tuple) + bodyList = [disc] + + # Generalized forces: Gravity: F_o; Auxiliary forces: F_p. + F_o = (O, F_O) + F_p = (P, F_P) + forceList = [F_o, F_p] + + # KanesMethod. + kane = KanesMethod( + N, q_ind= q[:3], u_ind= u[:3], kd_eqs=kindiffs, + q_dependent=q[3:], configuration_constraints = f_c, + u_dependent=u[3:], velocity_constraints= f_v, + u_auxiliary=ua + ) + + # fr, frstar, frstar_steady and kdd(kinematic differential equations). + (fr, frstar)= kane.kanes_equations(bodyList, forceList) + frstar_steady = frstar.subs(ud_zero).subs(u_dep_dict).subs(steady_conditions)\ + .subs({q[3]: -r*cos(q[1])}).expand() + kdd = kane.kindiffdict() + + assert Matrix(Fr_c).expand() == fr.expand() + assert Matrix(Fr_star_c.subs(kdd)).expand() == frstar.expand() + # These Matrices have some Integer(0) and some Float(0). Running under + # SymEngine gives different types of zero. + assert (simplify(Matrix(Fr_star_steady).expand()).xreplace({0:0.0}) == + simplify(frstar_steady.expand()).xreplace({0:0.0})) + + syms_in_forcing = find_dynamicsymbols(kane.forcing) + for qdi in qd: + assert qdi not in syms_in_forcing + + +def test_non_central_inertia(): + # This tests that the calculation of Fr* does not depend the point + # about which the inertia of a rigid body is defined. This test solves + # exercises 8.12, 8.17 from Kane 1985. + + # Declare symbols + q1, q2, q3 = dynamicsymbols('q1:4') + q1d, q2d, q3d = dynamicsymbols('q1:4', level=1) + u1, u2, u3, u4, u5 = dynamicsymbols('u1:6') + u_prime, R, M, g, e, f, theta = symbols('u\' R, M, g, e, f, theta') + a, b, mA, mB, IA, J, K, t = symbols('a b mA mB IA J K t') + Q1, Q2, Q3 = symbols('Q1, Q2 Q3') + IA22, IA23, IA33 = symbols('IA22 IA23 IA33') + + # Reference Frames + F = ReferenceFrame('F') + P = F.orientnew('P', 'axis', [-theta, F.y]) + A = P.orientnew('A', 'axis', [q1, P.x]) + A.set_ang_vel(F, u1*A.x + u3*A.z) + # define frames for wheels + B = A.orientnew('B', 'axis', [q2, A.z]) + C = A.orientnew('C', 'axis', [q3, A.z]) + B.set_ang_vel(A, u4 * A.z) + C.set_ang_vel(A, u5 * A.z) + + # define points D, S*, Q on frame A and their velocities + pD = Point('D') + pD.set_vel(A, 0) + # u3 will not change v_D_F since wheels are still assumed to roll without slip. + pD.set_vel(F, u2 * A.y) + + pS_star = pD.locatenew('S*', e*A.y) + pQ = pD.locatenew('Q', f*A.y - R*A.x) + for p in [pS_star, pQ]: + p.v2pt_theory(pD, F, A) + + # masscenters of bodies A, B, C + pA_star = pD.locatenew('A*', a*A.y) + pB_star = pD.locatenew('B*', b*A.z) + pC_star = pD.locatenew('C*', -b*A.z) + for p in [pA_star, pB_star, pC_star]: + p.v2pt_theory(pD, F, A) + + # points of B, C touching the plane P + pB_hat = pB_star.locatenew('B^', -R*A.x) + pC_hat = pC_star.locatenew('C^', -R*A.x) + pB_hat.v2pt_theory(pB_star, F, B) + pC_hat.v2pt_theory(pC_star, F, C) + + # the velocities of B^, C^ are zero since B, C are assumed to roll without slip + kde = [q1d - u1, q2d - u4, q3d - u5] + vc = [dot(p.vel(F), A.y) for p in [pB_hat, pC_hat]] + + # inertias of bodies A, B, C + # IA22, IA23, IA33 are not specified in the problem statement, but are + # necessary to define an inertia object. Although the values of + # IA22, IA23, IA33 are not known in terms of the variables given in the + # problem statement, they do not appear in the general inertia terms. + inertia_A = inertia(A, IA, IA22, IA33, 0, IA23, 0) + inertia_B = inertia(B, K, K, J) + inertia_C = inertia(C, K, K, J) + + # define the rigid bodies A, B, C + rbA = RigidBody('rbA', pA_star, A, mA, (inertia_A, pA_star)) + rbB = RigidBody('rbB', pB_star, B, mB, (inertia_B, pB_star)) + rbC = RigidBody('rbC', pC_star, C, mB, (inertia_C, pC_star)) + + km = KanesMethod(F, q_ind=[q1, q2, q3], u_ind=[u1, u2], kd_eqs=kde, + u_dependent=[u4, u5], velocity_constraints=vc, + u_auxiliary=[u3]) + + forces = [(pS_star, -M*g*F.x), (pQ, Q1*A.x + Q2*A.y + Q3*A.z)] + bodies = [rbA, rbB, rbC] + fr, fr_star = km.kanes_equations(bodies, forces) + vc_map = solve(vc, [u4, u5]) + + # KanesMethod returns the negative of Fr, Fr* as defined in Kane1985. + fr_star_expected = Matrix([ + -(IA + 2*J*b**2/R**2 + 2*K + + mA*a**2 + 2*mB*b**2) * u1.diff(t) - mA*a*u1*u2, + -(mA + 2*mB +2*J/R**2) * u2.diff(t) + mA*a*u1**2, + 0]) + t = trigsimp(fr_star.subs(vc_map).subs({u3: 0})).doit().expand() + assert ((fr_star_expected - t).expand() == zeros(3, 1)) + + # define inertias of rigid bodies A, B, C about point D + # I_S/O = I_S/S* + I_S*/O + bodies2 = [] + for rb, I_star in zip([rbA, rbB, rbC], [inertia_A, inertia_B, inertia_C]): + I = I_star + inertia_of_point_mass(rb.mass, + rb.masscenter.pos_from(pD), + rb.frame) + bodies2.append(RigidBody('', rb.masscenter, rb.frame, rb.mass, + (I, pD))) + fr2, fr_star2 = km.kanes_equations(bodies2, forces) + + t = trigsimp(fr_star2.subs(vc_map).subs({u3: 0})).doit() + assert (fr_star_expected - t).expand() == zeros(3, 1) + +def test_sub_qdot(): + # This test solves exercises 8.12, 8.17 from Kane 1985 and defines + # some velocities in terms of q, qdot. + + ## --- Declare symbols --- + q1, q2, q3 = dynamicsymbols('q1:4') + q1d, q2d, q3d = dynamicsymbols('q1:4', level=1) + u1, u2, u3 = dynamicsymbols('u1:4') + u_prime, R, M, g, e, f, theta = symbols('u\' R, M, g, e, f, theta') + a, b, mA, mB, IA, J, K, t = symbols('a b mA mB IA J K t') + IA22, IA23, IA33 = symbols('IA22 IA23 IA33') + Q1, Q2, Q3 = symbols('Q1 Q2 Q3') + + # --- Reference Frames --- + F = ReferenceFrame('F') + P = F.orientnew('P', 'axis', [-theta, F.y]) + A = P.orientnew('A', 'axis', [q1, P.x]) + A.set_ang_vel(F, u1*A.x + u3*A.z) + # define frames for wheels + B = A.orientnew('B', 'axis', [q2, A.z]) + C = A.orientnew('C', 'axis', [q3, A.z]) + + ## --- define points D, S*, Q on frame A and their velocities --- + pD = Point('D') + pD.set_vel(A, 0) + # u3 will not change v_D_F since wheels are still assumed to roll w/o slip + pD.set_vel(F, u2 * A.y) + + pS_star = pD.locatenew('S*', e*A.y) + pQ = pD.locatenew('Q', f*A.y - R*A.x) + # masscenters of bodies A, B, C + pA_star = pD.locatenew('A*', a*A.y) + pB_star = pD.locatenew('B*', b*A.z) + pC_star = pD.locatenew('C*', -b*A.z) + for p in [pS_star, pQ, pA_star, pB_star, pC_star]: + p.v2pt_theory(pD, F, A) + + # points of B, C touching the plane P + pB_hat = pB_star.locatenew('B^', -R*A.x) + pC_hat = pC_star.locatenew('C^', -R*A.x) + pB_hat.v2pt_theory(pB_star, F, B) + pC_hat.v2pt_theory(pC_star, F, C) + + # --- relate qdot, u --- + # the velocities of B^, C^ are zero since B, C are assumed to roll w/o slip + kde = [dot(p.vel(F), A.y) for p in [pB_hat, pC_hat]] + kde += [u1 - q1d] + kde_map = solve(kde, [q1d, q2d, q3d]) + for k, v in list(kde_map.items()): + kde_map[k.diff(t)] = v.diff(t) + + # inertias of bodies A, B, C + # IA22, IA23, IA33 are not specified in the problem statement, but are + # necessary to define an inertia object. Although the values of + # IA22, IA23, IA33 are not known in terms of the variables given in the + # problem statement, they do not appear in the general inertia terms. + inertia_A = inertia(A, IA, IA22, IA33, 0, IA23, 0) + inertia_B = inertia(B, K, K, J) + inertia_C = inertia(C, K, K, J) + + # define the rigid bodies A, B, C + rbA = RigidBody('rbA', pA_star, A, mA, (inertia_A, pA_star)) + rbB = RigidBody('rbB', pB_star, B, mB, (inertia_B, pB_star)) + rbC = RigidBody('rbC', pC_star, C, mB, (inertia_C, pC_star)) + + ## --- use kanes method --- + km = KanesMethod(F, [q1, q2, q3], [u1, u2], kd_eqs=kde, u_auxiliary=[u3]) + + forces = [(pS_star, -M*g*F.x), (pQ, Q1*A.x + Q2*A.y + Q3*A.z)] + bodies = [rbA, rbB, rbC] + + # Q2 = -u_prime * u2 * Q1 / sqrt(u2**2 + f**2 * u1**2) + # -u_prime * R * u2 / sqrt(u2**2 + f**2 * u1**2) = R / Q1 * Q2 + fr_expected = Matrix([ + f*Q3 + M*g*e*sin(theta)*cos(q1), + Q2 + M*g*sin(theta)*sin(q1), + e*M*g*cos(theta) - Q1*f - Q2*R]) + #Q1 * (f - u_prime * R * u2 / sqrt(u2**2 + f**2 * u1**2)))]) + fr_star_expected = Matrix([ + -(IA + 2*J*b**2/R**2 + 2*K + + mA*a**2 + 2*mB*b**2) * u1.diff(t) - mA*a*u1*u2, + -(mA + 2*mB +2*J/R**2) * u2.diff(t) + mA*a*u1**2, + 0]) + + fr, fr_star = km.kanes_equations(bodies, forces) + assert (fr.expand() == fr_expected.expand()) + assert ((fr_star_expected - trigsimp(fr_star)).expand() == zeros(3, 1)) + +def test_sub_qdot2(): + # This test solves exercises 8.3 from Kane 1985 and defines + # all velocities in terms of q, qdot. We check that the generalized active + # forces are correctly computed if u terms are only defined in the + # kinematic differential equations. + # + # This functionality was added in PR 8948. Without qdot/u substitution, the + # KanesMethod constructor will fail during the constraint initialization as + # the B matrix will be poorly formed and inversion of the dependent part + # will fail. + + g, m, Px, Py, Pz, R, t = symbols('g m Px Py Pz R t') + q = dynamicsymbols('q:5') + qd = dynamicsymbols('q:5', level=1) + u = dynamicsymbols('u:5') + + ## Define inertial, intermediate, and rigid body reference frames + A = ReferenceFrame('A') + B_prime = A.orientnew('B_prime', 'Axis', [q[0], A.z]) + B = B_prime.orientnew('B', 'Axis', [pi/2 - q[1], B_prime.x]) + C = B.orientnew('C', 'Axis', [q[2], B.z]) + + ## Define points of interest and their velocities + pO = Point('O') + pO.set_vel(A, 0) + + # R is the point in plane H that comes into contact with disk C. + pR = pO.locatenew('R', q[3]*A.x + q[4]*A.y) + pR.set_vel(A, pR.pos_from(pO).diff(t, A)) + pR.set_vel(B, 0) + + # C^ is the point in disk C that comes into contact with plane H. + pC_hat = pR.locatenew('C^', 0) + pC_hat.set_vel(C, 0) + + # C* is the point at the center of disk C. + pCs = pC_hat.locatenew('C*', R*B.y) + pCs.set_vel(C, 0) + pCs.set_vel(B, 0) + + # calculate velocites of points C* and C^ in frame A + pCs.v2pt_theory(pR, A, B) # points C* and R are fixed in frame B + pC_hat.v2pt_theory(pCs, A, C) # points C* and C^ are fixed in frame C + + ## Define forces on each point of the system + R_C_hat = Px*A.x + Py*A.y + Pz*A.z + R_Cs = -m*g*A.z + forces = [(pC_hat, R_C_hat), (pCs, R_Cs)] + + ## Define kinematic differential equations + # let ui = omega_C_A & bi (i = 1, 2, 3) + # u4 = qd4, u5 = qd5 + u_expr = [C.ang_vel_in(A) & uv for uv in B] + u_expr += qd[3:] + kde = [ui - e for ui, e in zip(u, u_expr)] + km1 = KanesMethod(A, q, u, kde) + fr1, _ = km1.kanes_equations([], forces) + + ## Calculate generalized active forces if we impose the condition that the + # disk C is rolling without slipping + u_indep = u[:3] + u_dep = list(set(u) - set(u_indep)) + vc = [pC_hat.vel(A) & uv for uv in [A.x, A.y]] + km2 = KanesMethod(A, q, u_indep, kde, + u_dependent=u_dep, velocity_constraints=vc) + fr2, _ = km2.kanes_equations([], forces) + + fr1_expected = Matrix([ + -R*g*m*sin(q[1]), + -R*(Px*cos(q[0]) + Py*sin(q[0]))*tan(q[1]), + R*(Px*cos(q[0]) + Py*sin(q[0])), + Px, + Py]) + fr2_expected = Matrix([ + -R*g*m*sin(q[1]), + 0, + 0]) + assert (trigsimp(fr1.expand()) == trigsimp(fr1_expected.expand())) + assert (trigsimp(fr2.expand()) == trigsimp(fr2_expected.expand())) diff --git a/lib/python3.12/site-packages/sympy/physics/mechanics/tests/test_kane3.py b/lib/python3.12/site-packages/sympy/physics/mechanics/tests/test_kane3.py new file mode 100644 index 0000000000000000000000000000000000000000..438759451cfb142c488b9b5c67ac269b668cac68 --- /dev/null +++ b/lib/python3.12/site-packages/sympy/physics/mechanics/tests/test_kane3.py @@ -0,0 +1,315 @@ +from sympy.core.numbers import pi +from sympy.core.symbol import symbols +from sympy.functions.elementary.miscellaneous import sqrt +from sympy.functions.elementary.trigonometric import acos, sin, cos +from sympy.matrices.dense import Matrix +from sympy.physics.mechanics import (ReferenceFrame, dynamicsymbols, + KanesMethod, inertia, Point, RigidBody, + dot) +from sympy.testing.pytest import slow + + +@slow +def test_bicycle(): + # Code to get equations of motion for a bicycle modeled as in: + # J.P Meijaard, Jim M Papadopoulos, Andy Ruina and A.L Schwab. Linearized + # dynamics equations for the balance and steer of a bicycle: a benchmark + # and review. Proceedings of The Royal Society (2007) 463, 1955-1982 + # doi: 10.1098/rspa.2007.1857 + + # Note that this code has been crudely ported from Autolev, which is the + # reason for some of the unusual naming conventions. It was purposefully as + # similar as possible in order to aide debugging. + + # Declare Coordinates & Speeds + # Simple definitions for qdots - qd = u + # Speeds are: + # - u1: yaw frame ang. rate + # - u2: roll frame ang. rate + # - u3: rear wheel frame ang. rate (spinning motion) + # - u4: frame ang. rate (pitching motion) + # - u5: steering frame ang. rate + # - u6: front wheel ang. rate (spinning motion) + # Wheel positions are ignorable coordinates, so they are not introduced. + q1, q2, q4, q5 = dynamicsymbols('q1 q2 q4 q5') + q1d, q2d, q4d, q5d = dynamicsymbols('q1 q2 q4 q5', 1) + u1, u2, u3, u4, u5, u6 = dynamicsymbols('u1 u2 u3 u4 u5 u6') + u1d, u2d, u3d, u4d, u5d, u6d = dynamicsymbols('u1 u2 u3 u4 u5 u6', 1) + + # Declare System's Parameters + WFrad, WRrad, htangle, forkoffset = symbols('WFrad WRrad htangle forkoffset') + forklength, framelength, forkcg1 = symbols('forklength framelength forkcg1') + forkcg3, framecg1, framecg3, Iwr11 = symbols('forkcg3 framecg1 framecg3 Iwr11') + Iwr22, Iwf11, Iwf22, Iframe11 = symbols('Iwr22 Iwf11 Iwf22 Iframe11') + Iframe22, Iframe33, Iframe31, Ifork11 = symbols('Iframe22 Iframe33 Iframe31 Ifork11') + Ifork22, Ifork33, Ifork31, g = symbols('Ifork22 Ifork33 Ifork31 g') + mframe, mfork, mwf, mwr = symbols('mframe mfork mwf mwr') + + # Set up reference frames for the system + # N - inertial + # Y - yaw + # R - roll + # WR - rear wheel, rotation angle is ignorable coordinate so not oriented + # Frame - bicycle frame + # TempFrame - statically rotated frame for easier reference inertia definition + # Fork - bicycle fork + # TempFork - statically rotated frame for easier reference inertia definition + # WF - front wheel, again posses a ignorable coordinate + N = ReferenceFrame('N') + Y = N.orientnew('Y', 'Axis', [q1, N.z]) + R = Y.orientnew('R', 'Axis', [q2, Y.x]) + Frame = R.orientnew('Frame', 'Axis', [q4 + htangle, R.y]) + WR = ReferenceFrame('WR') + TempFrame = Frame.orientnew('TempFrame', 'Axis', [-htangle, Frame.y]) + Fork = Frame.orientnew('Fork', 'Axis', [q5, Frame.x]) + TempFork = Fork.orientnew('TempFork', 'Axis', [-htangle, Fork.y]) + WF = ReferenceFrame('WF') + + # Kinematics of the Bicycle First block of code is forming the positions of + # the relevant points + # rear wheel contact -> rear wheel mass center -> frame mass center + + # frame/fork connection -> fork mass center + front wheel mass center -> + # front wheel contact point + WR_cont = Point('WR_cont') + WR_mc = WR_cont.locatenew('WR_mc', WRrad * R.z) + Steer = WR_mc.locatenew('Steer', framelength * Frame.z) + Frame_mc = WR_mc.locatenew('Frame_mc', - framecg1 * Frame.x + + framecg3 * Frame.z) + Fork_mc = Steer.locatenew('Fork_mc', - forkcg1 * Fork.x + + forkcg3 * Fork.z) + WF_mc = Steer.locatenew('WF_mc', forklength * Fork.x + forkoffset * Fork.z) + WF_cont = WF_mc.locatenew('WF_cont', WFrad * (dot(Fork.y, Y.z) * Fork.y - + Y.z).normalize()) + + # Set the angular velocity of each frame. + # Angular accelerations end up being calculated automatically by + # differentiating the angular velocities when first needed. + # u1 is yaw rate + # u2 is roll rate + # u3 is rear wheel rate + # u4 is frame pitch rate + # u5 is fork steer rate + # u6 is front wheel rate + Y.set_ang_vel(N, u1 * Y.z) + R.set_ang_vel(Y, u2 * R.x) + WR.set_ang_vel(Frame, u3 * Frame.y) + Frame.set_ang_vel(R, u4 * Frame.y) + Fork.set_ang_vel(Frame, u5 * Fork.x) + WF.set_ang_vel(Fork, u6 * Fork.y) + + # Form the velocities of the previously defined points, using the 2 - point + # theorem (written out by hand here). Accelerations again are calculated + # automatically when first needed. + WR_cont.set_vel(N, 0) + WR_mc.v2pt_theory(WR_cont, N, WR) + Steer.v2pt_theory(WR_mc, N, Frame) + Frame_mc.v2pt_theory(WR_mc, N, Frame) + Fork_mc.v2pt_theory(Steer, N, Fork) + WF_mc.v2pt_theory(Steer, N, Fork) + WF_cont.v2pt_theory(WF_mc, N, WF) + + # Sets the inertias of each body. Uses the inertia frame to construct the + # inertia dyadics. Wheel inertias are only defined by principle moments of + # inertia, and are in fact constant in the frame and fork reference frames; + # it is for this reason that the orientations of the wheels does not need + # to be defined. The frame and fork inertias are defined in the 'Temp' + # frames which are fixed to the appropriate body frames; this is to allow + # easier input of the reference values of the benchmark paper. Note that + # due to slightly different orientations, the products of inertia need to + # have their signs flipped; this is done later when entering the numerical + # value. + + Frame_I = (inertia(TempFrame, Iframe11, Iframe22, Iframe33, 0, 0, Iframe31), Frame_mc) + Fork_I = (inertia(TempFork, Ifork11, Ifork22, Ifork33, 0, 0, Ifork31), Fork_mc) + WR_I = (inertia(Frame, Iwr11, Iwr22, Iwr11), WR_mc) + WF_I = (inertia(Fork, Iwf11, Iwf22, Iwf11), WF_mc) + + # Declaration of the RigidBody containers. :: + + BodyFrame = RigidBody('BodyFrame', Frame_mc, Frame, mframe, Frame_I) + BodyFork = RigidBody('BodyFork', Fork_mc, Fork, mfork, Fork_I) + BodyWR = RigidBody('BodyWR', WR_mc, WR, mwr, WR_I) + BodyWF = RigidBody('BodyWF', WF_mc, WF, mwf, WF_I) + + # The kinematic differential equations; they are defined quite simply. Each + # entry in this list is equal to zero. + kd = [q1d - u1, q2d - u2, q4d - u4, q5d - u5] + + # The nonholonomic constraints are the velocity of the front wheel contact + # point dotted into the X, Y, and Z directions; the yaw frame is used as it + # is "closer" to the front wheel (1 less DCM connecting them). These + # constraints force the velocity of the front wheel contact point to be 0 + # in the inertial frame; the X and Y direction constraints enforce a + # "no-slip" condition, and the Z direction constraint forces the front + # wheel contact point to not move away from the ground frame, essentially + # replicating the holonomic constraint which does not allow the frame pitch + # to change in an invalid fashion. + + conlist_speed = [WF_cont.vel(N) & Y.x, WF_cont.vel(N) & Y.y, WF_cont.vel(N) & Y.z] + + # The holonomic constraint is that the position from the rear wheel contact + # point to the front wheel contact point when dotted into the + # normal-to-ground plane direction must be zero; effectively that the front + # and rear wheel contact points are always touching the ground plane. This + # is actually not part of the dynamic equations, but instead is necessary + # for the lineraization process. + + conlist_coord = [WF_cont.pos_from(WR_cont) & Y.z] + + # The force list; each body has the appropriate gravitational force applied + # at its mass center. + FL = [(Frame_mc, -mframe * g * Y.z), + (Fork_mc, -mfork * g * Y.z), + (WF_mc, -mwf * g * Y.z), + (WR_mc, -mwr * g * Y.z)] + BL = [BodyFrame, BodyFork, BodyWR, BodyWF] + + + # The N frame is the inertial frame, coordinates are supplied in the order + # of independent, dependent coordinates, as are the speeds. The kinematic + # differential equation are also entered here. Here the dependent speeds + # are specified, in the same order they were provided in earlier, along + # with the non-holonomic constraints. The dependent coordinate is also + # provided, with the holonomic constraint. Again, this is only provided + # for the linearization process. + + KM = KanesMethod(N, q_ind=[q1, q2, q5], + q_dependent=[q4], configuration_constraints=conlist_coord, + u_ind=[u2, u3, u5], + u_dependent=[u1, u4, u6], velocity_constraints=conlist_speed, + kd_eqs=kd, + constraint_solver="CRAMER") + (fr, frstar) = KM.kanes_equations(BL, FL) + + # This is the start of entering in the numerical values from the benchmark + # paper to validate the eigen values of the linearized equations from this + # model to the reference eigen values. Look at the aforementioned paper for + # more information. Some of these are intermediate values, used to + # transform values from the paper into the coordinate systems used in this + # model. + PaperRadRear = 0.3 + PaperRadFront = 0.35 + HTA = (pi / 2 - pi / 10).evalf() + TrailPaper = 0.08 + rake = (-(TrailPaper*sin(HTA)-(PaperRadFront*cos(HTA)))).evalf() + PaperWb = 1.02 + PaperFrameCgX = 0.3 + PaperFrameCgZ = 0.9 + PaperForkCgX = 0.9 + PaperForkCgZ = 0.7 + FrameLength = (PaperWb*sin(HTA)-(rake-(PaperRadFront-PaperRadRear)*cos(HTA))).evalf() + FrameCGNorm = ((PaperFrameCgZ - PaperRadRear-(PaperFrameCgX/sin(HTA))*cos(HTA))*sin(HTA)).evalf() + FrameCGPar = (PaperFrameCgX / sin(HTA) + (PaperFrameCgZ - PaperRadRear - PaperFrameCgX / sin(HTA) * cos(HTA)) * cos(HTA)).evalf() + tempa = (PaperForkCgZ - PaperRadFront) + tempb = (PaperWb-PaperForkCgX) + tempc = (sqrt(tempa**2+tempb**2)).evalf() + PaperForkL = (PaperWb*cos(HTA)-(PaperRadFront-PaperRadRear)*sin(HTA)).evalf() + ForkCGNorm = (rake+(tempc * sin(pi/2-HTA-acos(tempa/tempc)))).evalf() + ForkCGPar = (tempc * cos((pi/2-HTA)-acos(tempa/tempc))-PaperForkL).evalf() + + # Here is the final assembly of the numerical values. The symbol 'v' is the + # forward speed of the bicycle (a concept which only makes sense in the + # upright, static equilibrium case?). These are in a dictionary which will + # later be substituted in. Again the sign on the *product* of inertia + # values is flipped here, due to different orientations of coordinate + # systems. + v = symbols('v') + val_dict = {WFrad: PaperRadFront, + WRrad: PaperRadRear, + htangle: HTA, + forkoffset: rake, + forklength: PaperForkL, + framelength: FrameLength, + forkcg1: ForkCGPar, + forkcg3: ForkCGNorm, + framecg1: FrameCGNorm, + framecg3: FrameCGPar, + Iwr11: 0.0603, + Iwr22: 0.12, + Iwf11: 0.1405, + Iwf22: 0.28, + Ifork11: 0.05892, + Ifork22: 0.06, + Ifork33: 0.00708, + Ifork31: 0.00756, + Iframe11: 9.2, + Iframe22: 11, + Iframe33: 2.8, + Iframe31: -2.4, + mfork: 4, + mframe: 85, + mwf: 3, + mwr: 2, + g: 9.81, + q1: 0, + q2: 0, + q4: 0, + q5: 0, + u1: 0, + u2: 0, + u3: v / PaperRadRear, + u4: 0, + u5: 0, + u6: v / PaperRadFront} + + # Linearizes the forcing vector; the equations are set up as MM udot = + # forcing, where MM is the mass matrix, udot is the vector representing the + # time derivatives of the generalized speeds, and forcing is a vector which + # contains both external forcing terms and internal forcing terms, such as + # centripital or coriolis forces. This actually returns a matrix with as + # many rows as *total* coordinates and speeds, but only as many columns as + # independent coordinates and speeds. + + A, B, _ = KM.linearize( + A_and_B=True, + op_point={ + # Operating points for the accelerations are required for the + # linearizer to eliminate u' terms showing up in the coefficient + # matrices. + u1.diff(): 0, + u2.diff(): 0, + u3.diff(): 0, + u4.diff(): 0, + u5.diff(): 0, + u6.diff(): 0, + u1: 0, + u2: 0, + u3: v / PaperRadRear, + u4: 0, + u5: 0, + u6: v / PaperRadFront, + q1: 0, + q2: 0, + q4: 0, + q5: 0, + }, + linear_solver="CRAMER", + ) + # As mentioned above, the size of the linearized forcing terms is expanded + # to include both q's and u's, so the mass matrix must have this done as + # well. This will likely be changed to be part of the linearized process, + # for future reference. + A_s = A.xreplace(val_dict) + B_s = B.xreplace(val_dict) + + A_s = A_s.evalf() + B_s = B_s.evalf() + + # Finally, we construct an "A" matrix for the form xdot = A x (x being the + # state vector, although in this case, the sizes are a little off). The + # following line extracts only the minimum entries required for eigenvalue + # analysis, which correspond to rows and columns for lean, steer, lean + # rate, and steer rate. + A = A_s.extract([1, 2, 3, 5], [1, 2, 3, 5]) + + # Precomputed for comparison + Res = Matrix([[ 0, 0, 1.0, 0], + [ 0, 0, 0, 1.0], + [9.48977444677355, -0.891197738059089*v**2 - 0.571523173729245, -0.105522449805691*v, -0.330515398992311*v], + [11.7194768719633, -1.97171508499972*v**2 + 30.9087533932407, 3.67680523332152*v, -3.08486552743311*v]]) + + # Actual eigenvalue comparison + eps = 1.e-12 + for i in range(6): + error = Res.subs(v, i) - A.subs(v, i) + assert all(abs(x) < eps for x in error) diff --git a/lib/python3.12/site-packages/sympy/physics/mechanics/tests/test_kane4.py b/lib/python3.12/site-packages/sympy/physics/mechanics/tests/test_kane4.py new file mode 100644 index 0000000000000000000000000000000000000000..a44dd2d407056ea36669268d478780fc581def51 --- /dev/null +++ b/lib/python3.12/site-packages/sympy/physics/mechanics/tests/test_kane4.py @@ -0,0 +1,115 @@ +from sympy import (cos, sin, Matrix, symbols) +from sympy.physics.mechanics import (dynamicsymbols, ReferenceFrame, Point, + KanesMethod, Particle) + +def test_replace_qdots_in_force(): + # Test PR 16700 "Replaces qdots with us in force-list in kanes.py" + # The new functionality allows one to specify forces in qdots which will + # automatically be replaced with u:s which are defined by the kde supplied + # to KanesMethod. The test case is the double pendulum with interacting + # forces in the example of chapter 4.7 "CONTRIBUTING INTERACTION FORCES" + # in Ref. [1]. Reference list at end test function. + + q1, q2 = dynamicsymbols('q1, q2') + qd1, qd2 = dynamicsymbols('q1, q2', level=1) + u1, u2 = dynamicsymbols('u1, u2') + + l, m = symbols('l, m') + + N = ReferenceFrame('N') # Inertial frame + A = N.orientnew('A', 'Axis', (q1, N.z)) # Rod A frame + B = A.orientnew('B', 'Axis', (q2, N.z)) # Rod B frame + + O = Point('O') # Origo + O.set_vel(N, 0) + + P = O.locatenew('P', ( l * A.x )) # Point @ end of rod A + P.v2pt_theory(O, N, A) + + Q = P.locatenew('Q', ( l * B.x )) # Point @ end of rod B + Q.v2pt_theory(P, N, B) + + Ap = Particle('Ap', P, m) + Bp = Particle('Bp', Q, m) + + # The forces are specified below. sigma is the torsional spring stiffness + # and delta is the viscous damping coefficient acting between the two + # bodies. Here, we specify the viscous damper as function of qdots prior + # forming the kde. In more complex systems it not might be obvious which + # kde is most efficient, why it is convenient to specify viscous forces in + # qdots independently of the kde. + sig, delta = symbols('sigma, delta') + Ta = (sig * q2 + delta * qd2) * N.z + forces = [(A, Ta), (B, -Ta)] + + # Try different kdes. + kde1 = [u1 - qd1, u2 - qd2] + kde2 = [u1 - qd1, u2 - (qd1 + qd2)] + + KM1 = KanesMethod(N, [q1, q2], [u1, u2], kd_eqs=kde1) + fr1, fstar1 = KM1.kanes_equations([Ap, Bp], forces) + + KM2 = KanesMethod(N, [q1, q2], [u1, u2], kd_eqs=kde2) + fr2, fstar2 = KM2.kanes_equations([Ap, Bp], forces) + + # Check EOM for KM2: + # Mass and force matrix from p.6 in Ref. [2] with added forces from + # example of chapter 4.7 in [1] and without gravity. + forcing_matrix_expected = Matrix( [ [ m * l**2 * sin(q2) * u2**2 + sig * q2 + + delta * (u2 - u1)], + [ m * l**2 * sin(q2) * -u1**2 - sig * q2 + - delta * (u2 - u1)] ] ) + mass_matrix_expected = Matrix( [ [ 2 * m * l**2, m * l**2 * cos(q2) ], + [ m * l**2 * cos(q2), m * l**2 ] ] ) + + assert (KM2.mass_matrix.expand() == mass_matrix_expected.expand()) + assert (KM2.forcing.expand() == forcing_matrix_expected.expand()) + + # Check fr1 with reference fr_expected from [1] with u:s instead of qdots. + fr1_expected = Matrix([ 0, -(sig*q2 + delta * u2) ]) + assert fr1.expand() == fr1_expected.expand() + + # Check fr2 + fr2_expected = Matrix([sig * q2 + delta * (u2 - u1), + - sig * q2 - delta * (u2 - u1)]) + assert fr2.expand() == fr2_expected.expand() + + # Specifying forces in u:s should stay the same: + Ta = (sig * q2 + delta * u2) * N.z + forces = [(A, Ta), (B, -Ta)] + KM1 = KanesMethod(N, [q1, q2], [u1, u2], kd_eqs=kde1) + fr1, fstar1 = KM1.kanes_equations([Ap, Bp], forces) + + assert fr1.expand() == fr1_expected.expand() + + Ta = (sig * q2 + delta * (u2-u1)) * N.z + forces = [(A, Ta), (B, -Ta)] + KM2 = KanesMethod(N, [q1, q2], [u1, u2], kd_eqs=kde2) + fr2, fstar2 = KM2.kanes_equations([Ap, Bp], forces) + + assert fr2.expand() == fr2_expected.expand() + + # Test if we have a qubic qdot force: + Ta = (sig * q2 + delta * qd2**3) * N.z + forces = [(A, Ta), (B, -Ta)] + + KM1 = KanesMethod(N, [q1, q2], [u1, u2], kd_eqs=kde1) + fr1, fstar1 = KM1.kanes_equations([Ap, Bp], forces) + + fr1_cubic_expected = Matrix([ 0, -(sig*q2 + delta * u2**3) ]) + + assert fr1.expand() == fr1_cubic_expected.expand() + + KM2 = KanesMethod(N, [q1, q2], [u1, u2], kd_eqs=kde2) + fr2, fstar2 = KM2.kanes_equations([Ap, Bp], forces) + + fr2_cubic_expected = Matrix([sig * q2 + delta * (u2 - u1)**3, + - sig * q2 - delta * (u2 - u1)**3]) + + assert fr2.expand() == fr2_cubic_expected.expand() + + # References: + # [1] T.R. Kane, D. a Levinson, Dynamics Theory and Applications, 2005. + # [2] Arun K Banerjee, Flexible Multibody Dynamics:Efficient Formulations + # and Applications, John Wiley and Sons, Ltd, 2016. + # doi:http://dx.doi.org/10.1002/9781119015635. diff --git a/lib/python3.12/site-packages/sympy/physics/mechanics/tests/test_kane5.py b/lib/python3.12/site-packages/sympy/physics/mechanics/tests/test_kane5.py new file mode 100644 index 0000000000000000000000000000000000000000..1d0f863e8fa0f46bcd8ae729a1a8852b702bdafa --- /dev/null +++ b/lib/python3.12/site-packages/sympy/physics/mechanics/tests/test_kane5.py @@ -0,0 +1,128 @@ +from sympy import (zeros, Matrix, symbols, lambdify, sqrt, pi, + simplify) +from sympy.physics.mechanics import (dynamicsymbols, cross, inertia, RigidBody, + ReferenceFrame, KanesMethod) + + +def _create_rolling_disc(): + # Define symbols and coordinates + t = dynamicsymbols._t + q1, q2, q3, q4, q5, u1, u2, u3, u4, u5 = dynamicsymbols('q1:6 u1:6') + g, r, m = symbols('g r m') + # Define bodies and frames + ground = RigidBody('ground') + disc = RigidBody('disk', mass=m) + disc.inertia = (m * r ** 2 / 4 * inertia(disc.frame, 1, 2, 1), + disc.masscenter) + ground.masscenter.set_vel(ground.frame, 0) + disc.masscenter.set_vel(disc.frame, 0) + int_frame = ReferenceFrame('int_frame') + # Orient frames + int_frame.orient_body_fixed(ground.frame, (q1, q2, 0), 'zxy') + disc.frame.orient_axis(int_frame, int_frame.y, q3) + g_w_d = disc.frame.ang_vel_in(ground.frame) + disc.frame.set_ang_vel(ground.frame, + u1 * disc.x + u2 * disc.y + u3 * disc.z) + # Define points + cp = ground.masscenter.locatenew('contact_point', + q4 * ground.x + q5 * ground.y) + cp.set_vel(ground.frame, u4 * ground.x + u5 * ground.y) + disc.masscenter.set_pos(cp, r * int_frame.z) + disc.masscenter.set_vel(ground.frame, cross( + disc.frame.ang_vel_in(ground.frame), disc.masscenter.pos_from(cp))) + # Define kinematic differential equations + kdes = [g_w_d.dot(disc.x) - u1, g_w_d.dot(disc.y) - u2, + g_w_d.dot(disc.z) - u3, q4.diff(t) - u4, q5.diff(t) - u5] + # Define nonholonomic constraints + v0 = cp.vel(ground.frame) + cross( + disc.frame.ang_vel_in(int_frame), cp.pos_from(disc.masscenter)) + fnh = [v0.dot(ground.x), v0.dot(ground.y)] + # Define loads + loads = [(disc.masscenter, -disc.mass * g * ground.z)] + bodies = [disc] + return { + 'frame': ground.frame, + 'q_ind': [q1, q2, q3, q4, q5], + 'u_ind': [u1, u2, u3], + 'u_dep': [u4, u5], + 'kdes': kdes, + 'fnh': fnh, + 'bodies': bodies, + 'loads': loads + } + + +def _verify_rolling_disc_numerically(kane, all_zero=False): + q, u, p = dynamicsymbols('q1:6'), dynamicsymbols('u1:6'), symbols('g r m') + eval_sys = lambdify((q, u, p), (kane.mass_matrix_full, kane.forcing_full), + cse=True) + solve_sys = lambda q, u, p: Matrix.LUsolve( + *(Matrix(mat) for mat in eval_sys(q, u, p))) + solve_u_dep = lambdify((q, u[:3], p), kane._Ars * Matrix(u[:3]), cse=True) + eps = 1e-10 + p_vals = (9.81, 0.26, 3.43) + # First numeric test + q_vals = (0.3, 0.1, 1.97, -0.35, 2.27) + u_vals = [-0.2, 1.3, 0.15] + u_vals.extend(solve_u_dep(q_vals, u_vals, p_vals)[:2, 0]) + expected = Matrix([ + 0.126603940595934, 0.215942571601660, 1.28736069604936, + 0.319764288376543, 0.0989146857254898, -0.925848952664489, + -0.0181350656532944, 2.91695398184589, -0.00992793421754526, + 0.0412861634829171]) + assert all(abs(x) < eps for x in + (solve_sys(q_vals, u_vals, p_vals) - expected)) + # Second numeric test + q_vals = (3.97, -0.28, 8.2, -0.35, 2.27) + u_vals = [-0.25, -2.2, 0.62] + u_vals.extend(solve_u_dep(q_vals, u_vals, p_vals)[:2, 0]) + expected = Matrix([ + 0.0259159090798597, 0.668041660387416, -2.19283799213811, + 0.385441810852219, 0.420109283790573, 1.45030568179066, + -0.0110924422400793, -8.35617840186040, -0.154098542632173, + -0.146102664410010]) + assert all(abs(x) < eps for x in + (solve_sys(q_vals, u_vals, p_vals) - expected)) + if all_zero: + q_vals = (0, 0, 0, 0, 0) + u_vals = (0, 0, 0, 0, 0) + assert solve_sys(q_vals, u_vals, p_vals) == zeros(10, 1) + + +def test_kane_rolling_disc_lu(): + props = _create_rolling_disc() + kane = KanesMethod(props['frame'], props['q_ind'], props['u_ind'], + props['kdes'], u_dependent=props['u_dep'], + velocity_constraints=props['fnh'], + bodies=props['bodies'], forcelist=props['loads'], + explicit_kinematics=False, constraint_solver='LU') + kane.kanes_equations() + _verify_rolling_disc_numerically(kane) + + +def test_kane_rolling_disc_kdes_callable(): + props = _create_rolling_disc() + kane = KanesMethod( + props['frame'], props['q_ind'], props['u_ind'], props['kdes'], + u_dependent=props['u_dep'], velocity_constraints=props['fnh'], + bodies=props['bodies'], forcelist=props['loads'], + explicit_kinematics=False, + kd_eqs_solver=lambda A, b: simplify(A.LUsolve(b))) + q, u, p = dynamicsymbols('q1:6'), dynamicsymbols('u1:6'), symbols('g r m') + qd = dynamicsymbols('q1:6', 1) + eval_kdes = lambdify((q, qd, u, p), tuple(kane.kindiffdict().items())) + eps = 1e-10 + # Test with only zeros. If 'LU' would be used this would result in nan. + p_vals = (9.81, 0.25, 3.5) + zero_vals = (0, 0, 0, 0, 0) + assert all(abs(qdi - fui) < eps for qdi, fui in + eval_kdes(zero_vals, zero_vals, zero_vals, p_vals)) + # Test with some arbitrary values + q_vals = tuple(map(float, (pi / 6, pi / 3, pi / 2, 0.42, 0.62))) + qd_vals = tuple(map(float, (4, 1 / 3, 4 - 2 * sqrt(3), + 0.25 * (2 * sqrt(3) - 3), + 0.25 * (2 - sqrt(3))))) + u_vals = tuple(map(float, (-2, 4, 1 / 3, 0.25 * (-3 + 2 * sqrt(3)), + 0.25 * (-sqrt(3) + 2)))) + assert all(abs(qdi - fui) < eps for qdi, fui in + eval_kdes(q_vals, qd_vals, u_vals, p_vals)) diff --git a/lib/python3.12/site-packages/sympy/physics/mechanics/tests/test_lagrange.py b/lib/python3.12/site-packages/sympy/physics/mechanics/tests/test_lagrange.py new file mode 100644 index 0000000000000000000000000000000000000000..81552bc7a4d0f6766dc46dcd47b7c7b1b0151b3f --- /dev/null +++ b/lib/python3.12/site-packages/sympy/physics/mechanics/tests/test_lagrange.py @@ -0,0 +1,247 @@ +from sympy.physics.mechanics import (dynamicsymbols, ReferenceFrame, Point, + RigidBody, LagrangesMethod, Particle, + inertia, Lagrangian) +from sympy.core.function import (Derivative, Function) +from sympy.core.numbers import pi +from sympy.core.symbol import symbols +from sympy.functions.elementary.trigonometric import (cos, sin, tan) +from sympy.matrices.dense import Matrix +from sympy.simplify.simplify import simplify +from sympy.testing.pytest import raises + + +def test_invalid_coordinates(): + # Simple pendulum, but use symbol instead of dynamicsymbol + l, m, g = symbols('l m g') + q = symbols('q') # Generalized coordinate + N, O = ReferenceFrame('N'), Point('O') + O.set_vel(N, 0) + P = Particle('P', Point('P'), m) + P.point.set_pos(O, l * (sin(q) * N.x - cos(q) * N.y)) + P.potential_energy = m * g * P.point.pos_from(O).dot(N.y) + L = Lagrangian(N, P) + raises(ValueError, lambda: LagrangesMethod(L, [q], bodies=P)) + + +def test_disc_on_an_incline_plane(): + # Disc rolling on an inclined plane + # First the generalized coordinates are created. The mass center of the + # disc is located from top vertex of the inclined plane by the generalized + # coordinate 'y'. The orientation of the disc is defined by the angle + # 'theta'. The mass of the disc is 'm' and its radius is 'R'. The length of + # the inclined path is 'l', the angle of inclination is 'alpha'. 'g' is the + # gravitational constant. + y, theta = dynamicsymbols('y theta') + yd, thetad = dynamicsymbols('y theta', 1) + m, g, R, l, alpha = symbols('m g R l alpha') + + # Next, we create the inertial reference frame 'N'. A reference frame 'A' + # is attached to the inclined plane. Finally a frame is created which is attached to the disk. + N = ReferenceFrame('N') + A = N.orientnew('A', 'Axis', [pi/2 - alpha, N.z]) + B = A.orientnew('B', 'Axis', [-theta, A.z]) + + # Creating the disc 'D'; we create the point that represents the mass + # center of the disc and set its velocity. The inertia dyadic of the disc + # is created. Finally, we create the disc. + Do = Point('Do') + Do.set_vel(N, yd * A.x) + I = m * R**2/2 * B.z | B.z + D = RigidBody('D', Do, B, m, (I, Do)) + + # To construct the Lagrangian, 'L', of the disc, we determine its kinetic + # and potential energies, T and U, respectively. L is defined as the + # difference between T and U. + D.potential_energy = m * g * (l - y) * sin(alpha) + L = Lagrangian(N, D) + + # We then create the list of generalized coordinates and constraint + # equations. The constraint arises due to the disc rolling without slip on + # on the inclined path. We then invoke the 'LagrangesMethod' class and + # supply it the necessary arguments and generate the equations of motion. + # The'rhs' method solves for the q_double_dots (i.e. the second derivative + # with respect to time of the generalized coordinates and the lagrange + # multipliers. + q = [y, theta] + hol_coneqs = [y - R * theta] + m = LagrangesMethod(L, q, hol_coneqs=hol_coneqs) + m.form_lagranges_equations() + rhs = m.rhs() + rhs.simplify() + assert rhs[2] == 2*g*sin(alpha)/3 + + +def test_simp_pen(): + # This tests that the equations generated by LagrangesMethod are identical + # to those obtained by hand calculations. The system under consideration is + # the simple pendulum. + # We begin by creating the generalized coordinates as per the requirements + # of LagrangesMethod. Also we created the associate symbols + # that characterize the system: 'm' is the mass of the bob, l is the length + # of the massless rigid rod connecting the bob to a point O fixed in the + # inertial frame. + q, u = dynamicsymbols('q u') + qd, ud = dynamicsymbols('q u ', 1) + l, m, g = symbols('l m g') + + # We then create the inertial frame and a frame attached to the massless + # string following which we define the inertial angular velocity of the + # string. + N = ReferenceFrame('N') + A = N.orientnew('A', 'Axis', [q, N.z]) + A.set_ang_vel(N, qd * N.z) + + # Next, we create the point O and fix it in the inertial frame. We then + # locate the point P to which the bob is attached. Its corresponding + # velocity is then determined by the 'two point formula'. + O = Point('O') + O.set_vel(N, 0) + P = O.locatenew('P', l * A.x) + P.v2pt_theory(O, N, A) + + # The 'Particle' which represents the bob is then created and its + # Lagrangian generated. + Pa = Particle('Pa', P, m) + Pa.potential_energy = - m * g * l * cos(q) + L = Lagrangian(N, Pa) + + # The 'LagrangesMethod' class is invoked to obtain equations of motion. + lm = LagrangesMethod(L, [q]) + lm.form_lagranges_equations() + RHS = lm.rhs() + assert RHS[1] == -g*sin(q)/l + + +def test_nonminimal_pendulum(): + q1, q2 = dynamicsymbols('q1:3') + q1d, q2d = dynamicsymbols('q1:3', level=1) + L, m, t = symbols('L, m, t') + g = 9.8 + # Compose World Frame + N = ReferenceFrame('N') + pN = Point('N*') + pN.set_vel(N, 0) + # Create point P, the pendulum mass + P = pN.locatenew('P1', q1*N.x + q2*N.y) + P.set_vel(N, P.pos_from(pN).dt(N)) + pP = Particle('pP', P, m) + # Constraint Equations + f_c = Matrix([q1**2 + q2**2 - L**2]) + # Calculate the lagrangian, and form the equations of motion + Lag = Lagrangian(N, pP) + LM = LagrangesMethod(Lag, [q1, q2], hol_coneqs=f_c, + forcelist=[(P, m*g*N.x)], frame=N) + LM.form_lagranges_equations() + # Check solution + lam1 = LM.lam_vec[0, 0] + eom_sol = Matrix([[m*Derivative(q1, t, t) - 9.8*m + 2*lam1*q1], + [m*Derivative(q2, t, t) + 2*lam1*q2]]) + assert LM.eom == eom_sol + # Check multiplier solution + lam_sol = Matrix([(19.6*q1 + 2*q1d**2 + 2*q2d**2)/(4*q1**2/m + 4*q2**2/m)]) + assert simplify(LM.solve_multipliers(sol_type='Matrix')) == simplify(lam_sol) + + +def test_dub_pen(): + + # The system considered is the double pendulum. Like in the + # test of the simple pendulum above, we begin by creating the generalized + # coordinates and the simple generalized speeds and accelerations which + # will be used later. Following this we create frames and points necessary + # for the kinematics. The procedure isn't explicitly explained as this is + # similar to the simple pendulum. Also this is documented on the pydy.org + # website. + q1, q2 = dynamicsymbols('q1 q2') + q1d, q2d = dynamicsymbols('q1 q2', 1) + q1dd, q2dd = dynamicsymbols('q1 q2', 2) + u1, u2 = dynamicsymbols('u1 u2') + u1d, u2d = dynamicsymbols('u1 u2', 1) + l, m, g = symbols('l m g') + + N = ReferenceFrame('N') + A = N.orientnew('A', 'Axis', [q1, N.z]) + B = N.orientnew('B', 'Axis', [q2, N.z]) + + A.set_ang_vel(N, q1d * A.z) + B.set_ang_vel(N, q2d * A.z) + + O = Point('O') + P = O.locatenew('P', l * A.x) + R = P.locatenew('R', l * B.x) + + O.set_vel(N, 0) + P.v2pt_theory(O, N, A) + R.v2pt_theory(P, N, B) + + ParP = Particle('ParP', P, m) + ParR = Particle('ParR', R, m) + + ParP.potential_energy = - m * g * l * cos(q1) + ParR.potential_energy = - m * g * l * cos(q1) - m * g * l * cos(q2) + L = Lagrangian(N, ParP, ParR) + lm = LagrangesMethod(L, [q1, q2], bodies=[ParP, ParR]) + lm.form_lagranges_equations() + + assert simplify(l*m*(2*g*sin(q1) + l*sin(q1)*sin(q2)*q2dd + + l*sin(q1)*cos(q2)*q2d**2 - l*sin(q2)*cos(q1)*q2d**2 + + l*cos(q1)*cos(q2)*q2dd + 2*l*q1dd) - lm.eom[0]) == 0 + assert simplify(l*m*(g*sin(q2) + l*sin(q1)*sin(q2)*q1dd + - l*sin(q1)*cos(q2)*q1d**2 + l*sin(q2)*cos(q1)*q1d**2 + + l*cos(q1)*cos(q2)*q1dd + l*q2dd) - lm.eom[1]) == 0 + assert lm.bodies == [ParP, ParR] + + +def test_rolling_disc(): + # Rolling Disc Example + # Here the rolling disc is formed from the contact point up, removing the + # need to introduce generalized speeds. Only 3 configuration and 3 + # speed variables are need to describe this system, along with the + # disc's mass and radius, and the local gravity. + q1, q2, q3 = dynamicsymbols('q1 q2 q3') + q1d, q2d, q3d = dynamicsymbols('q1 q2 q3', 1) + r, m, g = symbols('r m g') + + # The kinematics are formed by a series of simple rotations. Each simple + # rotation creates a new frame, and the next rotation is defined by the new + # frame's basis vectors. This example uses a 3-1-2 series of rotations, or + # Z, X, Y series of rotations. Angular velocity for this is defined using + # the second frame's basis (the lean frame). + N = ReferenceFrame('N') + Y = N.orientnew('Y', 'Axis', [q1, N.z]) + L = Y.orientnew('L', 'Axis', [q2, Y.x]) + R = L.orientnew('R', 'Axis', [q3, L.y]) + + # This is the translational kinematics. We create a point with no velocity + # in N; this is the contact point between the disc and ground. Next we form + # the position vector from the contact point to the disc's center of mass. + # Finally we form the velocity and acceleration of the disc. + C = Point('C') + C.set_vel(N, 0) + Dmc = C.locatenew('Dmc', r * L.z) + Dmc.v2pt_theory(C, N, R) + + # Forming the inertia dyadic. + I = inertia(L, m/4 * r**2, m/2 * r**2, m/4 * r**2) + BodyD = RigidBody('BodyD', Dmc, R, m, (I, Dmc)) + + # Finally we form the equations of motion, using the same steps we did + # before. Supply the Lagrangian, the generalized speeds. + BodyD.potential_energy = - m * g * r * cos(q2) + Lag = Lagrangian(N, BodyD) + q = [q1, q2, q3] + q1 = Function('q1') + q2 = Function('q2') + q3 = Function('q3') + l = LagrangesMethod(Lag, q) + l.form_lagranges_equations() + RHS = l.rhs() + RHS.simplify() + t = symbols('t') + + assert (l.mass_matrix[3:6] == [0, 5*m*r**2/4, 0]) + assert RHS[4].simplify() == ( + (-8*g*sin(q2(t)) + r*(5*sin(2*q2(t))*Derivative(q1(t), t) + + 12*cos(q2(t))*Derivative(q3(t), t))*Derivative(q1(t), t))/(10*r)) + assert RHS[5] == (-5*cos(q2(t))*Derivative(q1(t), t) + 6*tan(q2(t) + )*Derivative(q3(t), t) + 4*Derivative(q1(t), t)/cos(q2(t)) + )*Derivative(q2(t), t) diff --git a/lib/python3.12/site-packages/sympy/physics/mechanics/tests/test_lagrange2.py b/lib/python3.12/site-packages/sympy/physics/mechanics/tests/test_lagrange2.py new file mode 100644 index 0000000000000000000000000000000000000000..d1bd1493306bba3f812192231e69ab314fd596b5 --- /dev/null +++ b/lib/python3.12/site-packages/sympy/physics/mechanics/tests/test_lagrange2.py @@ -0,0 +1,46 @@ +from sympy import symbols +from sympy.physics.mechanics import dynamicsymbols +from sympy.physics.mechanics import ReferenceFrame, Point, Particle +from sympy.physics.mechanics import LagrangesMethod, Lagrangian + +### This test asserts that a system with more than one external forces +### is acurately formed with Lagrange method (see issue #8626) + +def test_lagrange_2forces(): + ### Equations for two damped springs in serie with two forces + + ### generalized coordinates + q1, q2 = dynamicsymbols('q1, q2') + ### generalized speeds + q1d, q2d = dynamicsymbols('q1, q2', 1) + + ### Mass, spring strength, friction coefficient + m, k, nu = symbols('m, k, nu') + + N = ReferenceFrame('N') + O = Point('O') + + ### Two points + P1 = O.locatenew('P1', q1 * N.x) + P1.set_vel(N, q1d * N.x) + P2 = O.locatenew('P1', q2 * N.x) + P2.set_vel(N, q2d * N.x) + + pP1 = Particle('pP1', P1, m) + pP1.potential_energy = k * q1**2 / 2 + + pP2 = Particle('pP2', P2, m) + pP2.potential_energy = k * (q1 - q2)**2 / 2 + + #### Friction forces + forcelist = [(P1, - nu * q1d * N.x), + (P2, - nu * q2d * N.x)] + lag = Lagrangian(N, pP1, pP2) + + l_method = LagrangesMethod(lag, (q1, q2), forcelist=forcelist, frame=N) + l_method.form_lagranges_equations() + + eq1 = l_method.eom[0] + assert eq1.diff(q1d) == nu + eq2 = l_method.eom[1] + assert eq2.diff(q2d) == nu diff --git a/lib/python3.12/site-packages/sympy/physics/mechanics/tests/test_linearize.py b/lib/python3.12/site-packages/sympy/physics/mechanics/tests/test_linearize.py new file mode 100644 index 0000000000000000000000000000000000000000..ec62b960b71d7fce5a5504478431ca23eb371fe0 --- /dev/null +++ b/lib/python3.12/site-packages/sympy/physics/mechanics/tests/test_linearize.py @@ -0,0 +1,372 @@ +from sympy import symbols, Matrix, cos, sin, atan, sqrt, Rational +from sympy.core.sympify import sympify +from sympy.simplify.simplify import simplify +from sympy.solvers.solvers import solve +from sympy.physics.mechanics import dynamicsymbols, ReferenceFrame, Point,\ + dot, cross, inertia, KanesMethod, Particle, RigidBody, Lagrangian,\ + LagrangesMethod +from sympy.testing.pytest import slow + + +@slow +def test_linearize_rolling_disc_kane(): + # Symbols for time and constant parameters + t, r, m, g, v = symbols('t r m g v') + + # Configuration variables and their time derivatives + q1, q2, q3, q4, q5, q6 = q = dynamicsymbols('q1:7') + q1d, q2d, q3d, q4d, q5d, q6d = qd = [qi.diff(t) for qi in q] + + # Generalized speeds and their time derivatives + u = dynamicsymbols('u:6') + u1, u2, u3, u4, u5, u6 = u = dynamicsymbols('u1:7') + u1d, u2d, u3d, u4d, u5d, u6d = [ui.diff(t) for ui in u] + + # Reference frames + N = ReferenceFrame('N') # Inertial frame + NO = Point('NO') # Inertial origin + A = N.orientnew('A', 'Axis', [q1, N.z]) # Yaw intermediate frame + B = A.orientnew('B', 'Axis', [q2, A.x]) # Lean intermediate frame + C = B.orientnew('C', 'Axis', [q3, B.y]) # Disc fixed frame + CO = NO.locatenew('CO', q4*N.x + q5*N.y + q6*N.z) # Disc center + + # Disc angular velocity in N expressed using time derivatives of coordinates + w_c_n_qd = C.ang_vel_in(N) + w_b_n_qd = B.ang_vel_in(N) + + # Inertial angular velocity and angular acceleration of disc fixed frame + C.set_ang_vel(N, u1*B.x + u2*B.y + u3*B.z) + + # Disc center velocity in N expressed using time derivatives of coordinates + v_co_n_qd = CO.pos_from(NO).dt(N) + + # Disc center velocity in N expressed using generalized speeds + CO.set_vel(N, u4*C.x + u5*C.y + u6*C.z) + + # Disc Ground Contact Point + P = CO.locatenew('P', r*B.z) + P.v2pt_theory(CO, N, C) + + # Configuration constraint + f_c = Matrix([q6 - dot(CO.pos_from(P), N.z)]) + + # Velocity level constraints + f_v = Matrix([dot(P.vel(N), uv) for uv in C]) + + # Kinematic differential equations + kindiffs = Matrix([dot(w_c_n_qd - C.ang_vel_in(N), uv) for uv in B] + + [dot(v_co_n_qd - CO.vel(N), uv) for uv in N]) + qdots = solve(kindiffs, qd) + + # Set angular velocity of remaining frames + B.set_ang_vel(N, w_b_n_qd.subs(qdots)) + C.set_ang_acc(N, C.ang_vel_in(N).dt(B) + cross(B.ang_vel_in(N), C.ang_vel_in(N))) + + # Active forces + F_CO = m*g*A.z + + # Create inertia dyadic of disc C about point CO + I = (m * r**2) / 4 + J = (m * r**2) / 2 + I_C_CO = inertia(C, I, J, I) + + Disc = RigidBody('Disc', CO, C, m, (I_C_CO, CO)) + BL = [Disc] + FL = [(CO, F_CO)] + KM = KanesMethod(N, [q1, q2, q3, q4, q5], [u1, u2, u3], kd_eqs=kindiffs, + q_dependent=[q6], configuration_constraints=f_c, + u_dependent=[u4, u5, u6], velocity_constraints=f_v) + (fr, fr_star) = KM.kanes_equations(BL, FL) + + # Test generalized form equations + linearizer = KM.to_linearizer() + assert linearizer.f_c == f_c + assert linearizer.f_v == f_v + assert linearizer.f_a == f_v.diff(t).subs(KM.kindiffdict()) + sol = solve(linearizer.f_0 + linearizer.f_1, qd) + for qi in qdots.keys(): + assert sol[qi] == qdots[qi] + assert simplify(linearizer.f_2 + linearizer.f_3 - fr - fr_star) == Matrix([0, 0, 0]) + + # Perform the linearization + # Precomputed operating point + q_op = {q6: -r*cos(q2)} + u_op = {u1: 0, + u2: sin(q2)*q1d + q3d, + u3: cos(q2)*q1d, + u4: -r*(sin(q2)*q1d + q3d)*cos(q3), + u5: 0, + u6: -r*(sin(q2)*q1d + q3d)*sin(q3)} + qd_op = {q2d: 0, + q4d: -r*(sin(q2)*q1d + q3d)*cos(q1), + q5d: -r*(sin(q2)*q1d + q3d)*sin(q1), + q6d: 0} + ud_op = {u1d: 4*g*sin(q2)/(5*r) + sin(2*q2)*q1d**2/2 + 6*cos(q2)*q1d*q3d/5, + u2d: 0, + u3d: 0, + u4d: r*(sin(q2)*sin(q3)*q1d*q3d + sin(q3)*q3d**2), + u5d: r*(4*g*sin(q2)/(5*r) + sin(2*q2)*q1d**2/2 + 6*cos(q2)*q1d*q3d/5), + u6d: -r*(sin(q2)*cos(q3)*q1d*q3d + cos(q3)*q3d**2)} + + A, B = linearizer.linearize(op_point=[q_op, u_op, qd_op, ud_op], A_and_B=True, simplify=True) + + upright_nominal = {q1d: 0, q2: 0, m: 1, r: 1, g: 1} + + # Precomputed solution + A_sol = Matrix([[0, 0, 0, 0, 0, 0, 0, 1], + [0, 0, 0, 0, 0, 1, 0, 0], + [0, 0, 0, 0, 0, 0, 1, 0], + [sin(q1)*q3d, 0, 0, 0, 0, -sin(q1), -cos(q1), 0], + [-cos(q1)*q3d, 0, 0, 0, 0, cos(q1), -sin(q1), 0], + [0, Rational(4, 5), 0, 0, 0, 0, 0, 6*q3d/5], + [0, 0, 0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, -2*q3d, 0, 0]]) + B_sol = Matrix([]) + + # Check that linearization is correct + assert A.subs(upright_nominal) == A_sol + assert B.subs(upright_nominal) == B_sol + + # Check eigenvalues at critical speed are all zero: + assert sympify(A.subs(upright_nominal).subs(q3d, 1/sqrt(3))).eigenvals() == {0: 8} + + # Check whether alternative solvers work + # symengine doesn't support method='GJ' + linearizer = KM.to_linearizer(linear_solver='GJ') + A, B = linearizer.linearize(op_point=[q_op, u_op, qd_op, ud_op], + A_and_B=True, simplify=True) + assert A.subs(upright_nominal) == A_sol + assert B.subs(upright_nominal) == B_sol + +def test_linearize_pendulum_kane_minimal(): + q1 = dynamicsymbols('q1') # angle of pendulum + u1 = dynamicsymbols('u1') # Angular velocity + q1d = dynamicsymbols('q1', 1) # Angular velocity + L, m, t = symbols('L, m, t') + g = 9.8 + + # Compose world frame + N = ReferenceFrame('N') + pN = Point('N*') + pN.set_vel(N, 0) + + # A.x is along the pendulum + A = N.orientnew('A', 'axis', [q1, N.z]) + A.set_ang_vel(N, u1*N.z) + + # Locate point P relative to the origin N* + P = pN.locatenew('P', L*A.x) + P.v2pt_theory(pN, N, A) + pP = Particle('pP', P, m) + + # Create Kinematic Differential Equations + kde = Matrix([q1d - u1]) + + # Input the force resultant at P + R = m*g*N.x + + # Solve for eom with kanes method + KM = KanesMethod(N, q_ind=[q1], u_ind=[u1], kd_eqs=kde) + (fr, frstar) = KM.kanes_equations([pP], [(P, R)]) + + # Linearize + A, B, inp_vec = KM.linearize(A_and_B=True, simplify=True) + + assert A == Matrix([[0, 1], [-9.8*cos(q1)/L, 0]]) + assert B == Matrix([]) + +def test_linearize_pendulum_kane_nonminimal(): + # Create generalized coordinates and speeds for this non-minimal realization + # q1, q2 = N.x and N.y coordinates of pendulum + # u1, u2 = N.x and N.y velocities of pendulum + q1, q2 = dynamicsymbols('q1:3') + q1d, q2d = dynamicsymbols('q1:3', level=1) + u1, u2 = dynamicsymbols('u1:3') + u1d, u2d = dynamicsymbols('u1:3', level=1) + L, m, t = symbols('L, m, t') + g = 9.8 + + # Compose world frame + N = ReferenceFrame('N') + pN = Point('N*') + pN.set_vel(N, 0) + + # A.x is along the pendulum + theta1 = atan(q2/q1) + A = N.orientnew('A', 'axis', [theta1, N.z]) + + # Locate the pendulum mass + P = pN.locatenew('P1', q1*N.x + q2*N.y) + pP = Particle('pP', P, m) + + # Calculate the kinematic differential equations + kde = Matrix([q1d - u1, + q2d - u2]) + dq_dict = solve(kde, [q1d, q2d]) + + # Set velocity of point P + P.set_vel(N, P.pos_from(pN).dt(N).subs(dq_dict)) + + # Configuration constraint is length of pendulum + f_c = Matrix([P.pos_from(pN).magnitude() - L]) + + # Velocity constraint is that the velocity in the A.x direction is + # always zero (the pendulum is never getting longer). + f_v = Matrix([P.vel(N).express(A).dot(A.x)]) + f_v.simplify() + + # Acceleration constraints is the time derivative of the velocity constraint + f_a = f_v.diff(t) + f_a.simplify() + + # Input the force resultant at P + R = m*g*N.x + + # Derive the equations of motion using the KanesMethod class. + KM = KanesMethod(N, q_ind=[q2], u_ind=[u2], q_dependent=[q1], + u_dependent=[u1], configuration_constraints=f_c, + velocity_constraints=f_v, acceleration_constraints=f_a, kd_eqs=kde) + (fr, frstar) = KM.kanes_equations([pP], [(P, R)]) + + # Set the operating point to be straight down, and non-moving + q_op = {q1: L, q2: 0} + u_op = {u1: 0, u2: 0} + ud_op = {u1d: 0, u2d: 0} + + A, B, inp_vec = KM.linearize(op_point=[q_op, u_op, ud_op], A_and_B=True, + simplify=True) + + assert A.expand() == Matrix([[0, 1], [-9.8/L, 0]]) + assert B == Matrix([]) + + + # symengine doesn't support method='GJ' + A, B, inp_vec = KM.linearize(op_point=[q_op, u_op, ud_op], A_and_B=True, + simplify=True, linear_solver='GJ') + + assert A.expand() == Matrix([[0, 1], [-9.8/L, 0]]) + assert B == Matrix([]) + + A, B, inp_vec = KM.linearize(op_point=[q_op, u_op, ud_op], + A_and_B=True, + simplify=True, + linear_solver=lambda A, b: A.LUsolve(b)) + + assert A.expand() == Matrix([[0, 1], [-9.8/L, 0]]) + assert B == Matrix([]) + + +def test_linearize_pendulum_lagrange_minimal(): + q1 = dynamicsymbols('q1') # angle of pendulum + q1d = dynamicsymbols('q1', 1) # Angular velocity + L, m, t = symbols('L, m, t') + g = 9.8 + + # Compose world frame + N = ReferenceFrame('N') + pN = Point('N*') + pN.set_vel(N, 0) + + # A.x is along the pendulum + A = N.orientnew('A', 'axis', [q1, N.z]) + A.set_ang_vel(N, q1d*N.z) + + # Locate point P relative to the origin N* + P = pN.locatenew('P', L*A.x) + P.v2pt_theory(pN, N, A) + pP = Particle('pP', P, m) + + # Solve for eom with Lagranges method + Lag = Lagrangian(N, pP) + LM = LagrangesMethod(Lag, [q1], forcelist=[(P, m*g*N.x)], frame=N) + LM.form_lagranges_equations() + + # Linearize + A, B, inp_vec = LM.linearize([q1], [q1d], A_and_B=True) + + assert simplify(A) == Matrix([[0, 1], [-9.8*cos(q1)/L, 0]]) + assert B == Matrix([]) + + # Check an alternative solver + A, B, inp_vec = LM.linearize([q1], [q1d], A_and_B=True, linear_solver='GJ') + + assert simplify(A) == Matrix([[0, 1], [-9.8*cos(q1)/L, 0]]) + assert B == Matrix([]) + + +def test_linearize_pendulum_lagrange_nonminimal(): + q1, q2 = dynamicsymbols('q1:3') + q1d, q2d = dynamicsymbols('q1:3', level=1) + L, m, t = symbols('L, m, t') + g = 9.8 + # Compose World Frame + N = ReferenceFrame('N') + pN = Point('N*') + pN.set_vel(N, 0) + # A.x is along the pendulum + theta1 = atan(q2/q1) + A = N.orientnew('A', 'axis', [theta1, N.z]) + # Create point P, the pendulum mass + P = pN.locatenew('P1', q1*N.x + q2*N.y) + P.set_vel(N, P.pos_from(pN).dt(N)) + pP = Particle('pP', P, m) + # Constraint Equations + f_c = Matrix([q1**2 + q2**2 - L**2]) + # Calculate the lagrangian, and form the equations of motion + Lag = Lagrangian(N, pP) + LM = LagrangesMethod(Lag, [q1, q2], hol_coneqs=f_c, forcelist=[(P, m*g*N.x)], frame=N) + LM.form_lagranges_equations() + # Compose operating point + op_point = {q1: L, q2: 0, q1d: 0, q2d: 0, q1d.diff(t): 0, q2d.diff(t): 0} + # Solve for multiplier operating point + lam_op = LM.solve_multipliers(op_point=op_point) + op_point.update(lam_op) + # Perform the Linearization + A, B, inp_vec = LM.linearize([q2], [q2d], [q1], [q1d], + op_point=op_point, A_and_B=True) + assert simplify(A) == Matrix([[0, 1], [-9.8/L, 0]]) + assert B == Matrix([]) + + # Check if passing a function to linear_solver works + A, B, inp_vec = LM.linearize([q2], [q2d], [q1], [q1d], op_point=op_point, + A_and_B=True, linear_solver=lambda A, b: + A.LUsolve(b)) + assert simplify(A) == Matrix([[0, 1], [-9.8/L, 0]]) + assert B == Matrix([]) + +def test_linearize_rolling_disc_lagrange(): + q1, q2, q3 = q = dynamicsymbols('q1 q2 q3') + q1d, q2d, q3d = qd = dynamicsymbols('q1 q2 q3', 1) + r, m, g = symbols('r m g') + + N = ReferenceFrame('N') + Y = N.orientnew('Y', 'Axis', [q1, N.z]) + L = Y.orientnew('L', 'Axis', [q2, Y.x]) + R = L.orientnew('R', 'Axis', [q3, L.y]) + + C = Point('C') + C.set_vel(N, 0) + Dmc = C.locatenew('Dmc', r * L.z) + Dmc.v2pt_theory(C, N, R) + + I = inertia(L, m / 4 * r**2, m / 2 * r**2, m / 4 * r**2) + BodyD = RigidBody('BodyD', Dmc, R, m, (I, Dmc)) + BodyD.potential_energy = - m * g * r * cos(q2) + + Lag = Lagrangian(N, BodyD) + l = LagrangesMethod(Lag, q) + l.form_lagranges_equations() + + # Linearize about steady-state upright rolling + op_point = {q1: 0, q2: 0, q3: 0, + q1d: 0, q2d: 0, + q1d.diff(): 0, q2d.diff(): 0, q3d.diff(): 0} + A = l.linearize(q_ind=q, qd_ind=qd, op_point=op_point, A_and_B=True)[0] + sol = Matrix([[0, 0, 0, 1, 0, 0], + [0, 0, 0, 0, 1, 0], + [0, 0, 0, 0, 0, 1], + [0, 0, 0, 0, -6*q3d, 0], + [0, -4*g/(5*r), 0, 6*q3d/5, 0, 0], + [0, 0, 0, 0, 0, 0]]) + + assert A == sol diff --git a/lib/python3.12/site-packages/sympy/physics/mechanics/tests/test_loads.py b/lib/python3.12/site-packages/sympy/physics/mechanics/tests/test_loads.py new file mode 100644 index 0000000000000000000000000000000000000000..8aa0cec14887f0778fc1e60e7ff33830ceef72d3 --- /dev/null +++ b/lib/python3.12/site-packages/sympy/physics/mechanics/tests/test_loads.py @@ -0,0 +1,86 @@ +from pytest import raises + +from sympy import symbols +from sympy.physics.mechanics import (RigidBody, Particle, ReferenceFrame, Point, + outer, dynamicsymbols, Force, Torque) +from sympy.physics.mechanics.loads import gravity, _parse_load + + +def test_force_default(): + N = ReferenceFrame('N') + Po = Point('Po') + f1 = Force(Po, N.x) + assert f1.point == Po + assert f1.force == N.x + assert f1.__repr__() == 'Force(point=Po, force=N.x)' + # Test tuple behaviour + assert isinstance(f1, tuple) + assert f1[0] == Po + assert f1[1] == N.x + assert f1 == (Po, N.x) + assert f1 != (N.x, Po) + assert f1 != (Po, N.x + N.y) + assert f1 != (Point('Co'), N.x) + # Test body as input + P = Particle('P', Po) + f2 = Force(P, N.x) + assert f1 == f2 + + +def test_torque_default(): + N = ReferenceFrame('N') + f1 = Torque(N, N.x) + assert f1.frame == N + assert f1.torque == N.x + assert f1.__repr__() == 'Torque(frame=N, torque=N.x)' + # Test tuple behaviour + assert isinstance(f1, tuple) + assert f1[0] == N + assert f1[1] == N.x + assert f1 == (N, N.x) + assert f1 != (N.x, N) + assert f1 != (N, N.x + N.y) + assert f1 != (ReferenceFrame('A'), N.x) + # Test body as input + rb = RigidBody('P', frame=N) + f2 = Torque(rb, N.x) + assert f1 == f2 + + +def test_gravity(): + N = ReferenceFrame('N') + m, M, g = symbols('m M g') + F1, F2 = dynamicsymbols('F1 F2') + po = Point('po') + pa = Particle('pa', po, m) + A = ReferenceFrame('A') + P = Point('P') + I = outer(A.x, A.x) + B = RigidBody('B', P, A, M, (I, P)) + forceList = [(po, F1), (P, F2)] + forceList.extend(gravity(g * N.y, pa, B)) + l = [(po, F1), (P, F2), (po, g * m * N.y), (P, g * M * N.y)] + + for i in range(len(l)): + for j in range(len(l[i])): + assert forceList[i][j] == l[i][j] + + +def test_parse_loads(): + N = ReferenceFrame('N') + po = Point('po') + assert _parse_load(Force(po, N.z)) == (po, N.z) + assert _parse_load(Torque(N, N.x)) == (N, N.x) + f1 = _parse_load((po, N.x)) # Test whether a force is recognized + assert isinstance(f1, Force) + assert f1 == Force(po, N.x) + t1 = _parse_load((N, N.y)) # Test whether a torque is recognized + assert isinstance(t1, Torque) + assert t1 == Torque(N, N.y) + # Bodies should be undetermined (even in case of a Particle) + raises(ValueError, lambda: _parse_load((Particle('pa', po), N.x))) + raises(ValueError, lambda: _parse_load((RigidBody('pa', po, N), N.x))) + # Invalid tuple length + raises(ValueError, lambda: _parse_load((po, N.x, po, N.x))) + # Invalid type + raises(TypeError, lambda: _parse_load([po, N.x])) diff --git a/lib/python3.12/site-packages/sympy/physics/mechanics/tests/test_method.py b/lib/python3.12/site-packages/sympy/physics/mechanics/tests/test_method.py new file mode 100644 index 0000000000000000000000000000000000000000..4a8fd5fb50c3178f5a5cdab1e80423df8b52f525 --- /dev/null +++ b/lib/python3.12/site-packages/sympy/physics/mechanics/tests/test_method.py @@ -0,0 +1,5 @@ +from sympy.physics.mechanics.method import _Methods +from sympy.testing.pytest import raises + +def test_method(): + raises(TypeError, lambda: _Methods()) diff --git a/lib/python3.12/site-packages/sympy/physics/mechanics/tests/test_models.py b/lib/python3.12/site-packages/sympy/physics/mechanics/tests/test_models.py new file mode 100644 index 0000000000000000000000000000000000000000..2b3d3ae89b44d774ead1a3ea641a8274ba951638 --- /dev/null +++ b/lib/python3.12/site-packages/sympy/physics/mechanics/tests/test_models.py @@ -0,0 +1,117 @@ +import sympy.physics.mechanics.models as models +from sympy import (cos, sin, Matrix, symbols, zeros) +from sympy.simplify.simplify import simplify +from sympy.physics.mechanics import (dynamicsymbols) + + +def test_multi_mass_spring_damper_inputs(): + + c0, k0, m0 = symbols("c0 k0 m0") + g = symbols("g") + v0, x0, f0 = dynamicsymbols("v0 x0 f0") + + kane1 = models.multi_mass_spring_damper(1) + massmatrix1 = Matrix([[m0]]) + forcing1 = Matrix([[-c0*v0 - k0*x0]]) + assert simplify(massmatrix1 - kane1.mass_matrix) == Matrix([0]) + assert simplify(forcing1 - kane1.forcing) == Matrix([0]) + + kane2 = models.multi_mass_spring_damper(1, True) + massmatrix2 = Matrix([[m0]]) + forcing2 = Matrix([[-c0*v0 + g*m0 - k0*x0]]) + assert simplify(massmatrix2 - kane2.mass_matrix) == Matrix([0]) + assert simplify(forcing2 - kane2.forcing) == Matrix([0]) + + kane3 = models.multi_mass_spring_damper(1, True, True) + massmatrix3 = Matrix([[m0]]) + forcing3 = Matrix([[-c0*v0 + g*m0 - k0*x0 + f0]]) + assert simplify(massmatrix3 - kane3.mass_matrix) == Matrix([0]) + assert simplify(forcing3 - kane3.forcing) == Matrix([0]) + + kane4 = models.multi_mass_spring_damper(1, False, True) + massmatrix4 = Matrix([[m0]]) + forcing4 = Matrix([[-c0*v0 - k0*x0 + f0]]) + assert simplify(massmatrix4 - kane4.mass_matrix) == Matrix([0]) + assert simplify(forcing4 - kane4.forcing) == Matrix([0]) + + +def test_multi_mass_spring_damper_higher_order(): + c0, k0, m0 = symbols("c0 k0 m0") + c1, k1, m1 = symbols("c1 k1 m1") + c2, k2, m2 = symbols("c2 k2 m2") + v0, x0 = dynamicsymbols("v0 x0") + v1, x1 = dynamicsymbols("v1 x1") + v2, x2 = dynamicsymbols("v2 x2") + + kane1 = models.multi_mass_spring_damper(3) + massmatrix1 = Matrix([[m0 + m1 + m2, m1 + m2, m2], + [m1 + m2, m1 + m2, m2], + [m2, m2, m2]]) + forcing1 = Matrix([[-c0*v0 - k0*x0], + [-c1*v1 - k1*x1], + [-c2*v2 - k2*x2]]) + assert simplify(massmatrix1 - kane1.mass_matrix) == zeros(3) + assert simplify(forcing1 - kane1.forcing) == Matrix([0, 0, 0]) + + +def test_n_link_pendulum_on_cart_inputs(): + l0, m0 = symbols("l0 m0") + m1 = symbols("m1") + g = symbols("g") + q0, q1, F, T1 = dynamicsymbols("q0 q1 F T1") + u0, u1 = dynamicsymbols("u0 u1") + + kane1 = models.n_link_pendulum_on_cart(1) + massmatrix1 = Matrix([[m0 + m1, -l0*m1*cos(q1)], + [-l0*m1*cos(q1), l0**2*m1]]) + forcing1 = Matrix([[-l0*m1*u1**2*sin(q1) + F], [g*l0*m1*sin(q1)]]) + assert simplify(massmatrix1 - kane1.mass_matrix) == zeros(2) + assert simplify(forcing1 - kane1.forcing) == Matrix([0, 0]) + + kane2 = models.n_link_pendulum_on_cart(1, False) + massmatrix2 = Matrix([[m0 + m1, -l0*m1*cos(q1)], + [-l0*m1*cos(q1), l0**2*m1]]) + forcing2 = Matrix([[-l0*m1*u1**2*sin(q1)], [g*l0*m1*sin(q1)]]) + assert simplify(massmatrix2 - kane2.mass_matrix) == zeros(2) + assert simplify(forcing2 - kane2.forcing) == Matrix([0, 0]) + + kane3 = models.n_link_pendulum_on_cart(1, False, True) + massmatrix3 = Matrix([[m0 + m1, -l0*m1*cos(q1)], + [-l0*m1*cos(q1), l0**2*m1]]) + forcing3 = Matrix([[-l0*m1*u1**2*sin(q1)], [g*l0*m1*sin(q1) + T1]]) + assert simplify(massmatrix3 - kane3.mass_matrix) == zeros(2) + assert simplify(forcing3 - kane3.forcing) == Matrix([0, 0]) + + kane4 = models.n_link_pendulum_on_cart(1, True, False) + massmatrix4 = Matrix([[m0 + m1, -l0*m1*cos(q1)], + [-l0*m1*cos(q1), l0**2*m1]]) + forcing4 = Matrix([[-l0*m1*u1**2*sin(q1) + F], [g*l0*m1*sin(q1)]]) + assert simplify(massmatrix4 - kane4.mass_matrix) == zeros(2) + assert simplify(forcing4 - kane4.forcing) == Matrix([0, 0]) + + +def test_n_link_pendulum_on_cart_higher_order(): + l0, m0 = symbols("l0 m0") + l1, m1 = symbols("l1 m1") + m2 = symbols("m2") + g = symbols("g") + q0, q1, q2 = dynamicsymbols("q0 q1 q2") + u0, u1, u2 = dynamicsymbols("u0 u1 u2") + F, T1 = dynamicsymbols("F T1") + + kane1 = models.n_link_pendulum_on_cart(2) + massmatrix1 = Matrix([[m0 + m1 + m2, -l0*m1*cos(q1) - l0*m2*cos(q1), + -l1*m2*cos(q2)], + [-l0*m1*cos(q1) - l0*m2*cos(q1), l0**2*m1 + l0**2*m2, + l0*l1*m2*(sin(q1)*sin(q2) + cos(q1)*cos(q2))], + [-l1*m2*cos(q2), + l0*l1*m2*(sin(q1)*sin(q2) + cos(q1)*cos(q2)), + l1**2*m2]]) + forcing1 = Matrix([[-l0*m1*u1**2*sin(q1) - l0*m2*u1**2*sin(q1) - + l1*m2*u2**2*sin(q2) + F], + [g*l0*m1*sin(q1) + g*l0*m2*sin(q1) - + l0*l1*m2*(sin(q1)*cos(q2) - sin(q2)*cos(q1))*u2**2], + [g*l1*m2*sin(q2) - l0*l1*m2*(-sin(q1)*cos(q2) + + sin(q2)*cos(q1))*u1**2]]) + assert simplify(massmatrix1 - kane1.mass_matrix) == zeros(3) + assert simplify(forcing1 - kane1.forcing) == Matrix([0, 0, 0]) diff --git a/lib/python3.12/site-packages/sympy/physics/mechanics/tests/test_particle.py b/lib/python3.12/site-packages/sympy/physics/mechanics/tests/test_particle.py new file mode 100644 index 0000000000000000000000000000000000000000..8eec80275b532055eacaf2339a276c0fd19b330a --- /dev/null +++ b/lib/python3.12/site-packages/sympy/physics/mechanics/tests/test_particle.py @@ -0,0 +1,78 @@ +from sympy import symbols +from sympy.physics.mechanics import Point, Particle, ReferenceFrame, inertia +from sympy.physics.mechanics.body_base import BodyBase +from sympy.testing.pytest import raises, warns_deprecated_sympy + + +def test_particle_default(): + # Test default + p = Particle('P') + assert p.name == 'P' + assert p.mass == symbols('P_mass') + assert p.masscenter.name == 'P_masscenter' + assert p.potential_energy == 0 + assert p.__str__() == 'P' + assert p.__repr__() == ("Particle('P', masscenter=P_masscenter, " + "mass=P_mass)") + raises(AttributeError, lambda: p.frame) + + +def test_particle(): + # Test initializing with parameters + m, m2, v1, v2, v3, r, g, h = symbols('m m2 v1 v2 v3 r g h') + P = Point('P') + P2 = Point('P2') + p = Particle('pa', P, m) + assert isinstance(p, BodyBase) + assert p.mass == m + assert p.point == P + # Test the mass setter + p.mass = m2 + assert p.mass == m2 + # Test the point setter + p.point = P2 + assert p.point == P2 + # Test the linear momentum function + N = ReferenceFrame('N') + O = Point('O') + P2.set_pos(O, r * N.y) + P2.set_vel(N, v1 * N.x) + raises(TypeError, lambda: Particle(P, P, m)) + raises(TypeError, lambda: Particle('pa', m, m)) + assert p.linear_momentum(N) == m2 * v1 * N.x + assert p.angular_momentum(O, N) == -m2 * r * v1 * N.z + P2.set_vel(N, v2 * N.y) + assert p.linear_momentum(N) == m2 * v2 * N.y + assert p.angular_momentum(O, N) == 0 + P2.set_vel(N, v3 * N.z) + assert p.linear_momentum(N) == m2 * v3 * N.z + assert p.angular_momentum(O, N) == m2 * r * v3 * N.x + P2.set_vel(N, v1 * N.x + v2 * N.y + v3 * N.z) + assert p.linear_momentum(N) == m2 * (v1 * N.x + v2 * N.y + v3 * N.z) + assert p.angular_momentum(O, N) == m2 * r * (v3 * N.x - v1 * N.z) + p.potential_energy = m * g * h + assert p.potential_energy == m * g * h + # TODO make the result not be system-dependent + assert p.kinetic_energy( + N) in [m2 * (v1 ** 2 + v2 ** 2 + v3 ** 2) / 2, + m2 * v1 ** 2 / 2 + m2 * v2 ** 2 / 2 + m2 * v3 ** 2 / 2] + + +def test_parallel_axis(): + N = ReferenceFrame('N') + m, a, b = symbols('m, a, b') + o = Point('o') + p = o.locatenew('p', a * N.x + b * N.y) + P = Particle('P', o, m) + Ip = P.parallel_axis(p, N) + Ip_expected = inertia(N, m * b ** 2, m * a ** 2, m * (a ** 2 + b ** 2), + ixy=-m * a * b) + assert Ip == Ip_expected + + +def test_deprecated_set_potential_energy(): + m, g, h = symbols('m g h') + P = Point('P') + p = Particle('pa', P, m) + with warns_deprecated_sympy(): + p.set_potential_energy(m * g * h) diff --git a/lib/python3.12/site-packages/sympy/physics/mechanics/tests/test_pathway.py b/lib/python3.12/site-packages/sympy/physics/mechanics/tests/test_pathway.py new file mode 100644 index 0000000000000000000000000000000000000000..49dc4bd4d61300745833f9d32f3a91d9054c4839 --- /dev/null +++ b/lib/python3.12/site-packages/sympy/physics/mechanics/tests/test_pathway.py @@ -0,0 +1,691 @@ +"""Tests for the ``sympy.physics.mechanics.pathway.py`` module.""" + +import pytest + +from sympy import ( + Rational, + Symbol, + cos, + pi, + sin, + sqrt, +) +from sympy.physics.mechanics import ( + Force, + LinearPathway, + ObstacleSetPathway, + PathwayBase, + Point, + ReferenceFrame, + WrappingCylinder, + WrappingGeometryBase, + WrappingPathway, + WrappingSphere, + dynamicsymbols, +) +from sympy.simplify.simplify import simplify + + +def _simplify_loads(loads): + return [ + load.__class__(load.location, load.vector.simplify()) + for load in loads + ] + + +class TestLinearPathway: + + def test_is_pathway_base_subclass(self): + assert issubclass(LinearPathway, PathwayBase) + + @staticmethod + @pytest.mark.parametrize( + 'args, kwargs', + [ + ((Point('pA'), Point('pB')), {}), + ] + ) + def test_valid_constructor(args, kwargs): + pointA, pointB = args + instance = LinearPathway(*args, **kwargs) + assert isinstance(instance, LinearPathway) + assert hasattr(instance, 'attachments') + assert len(instance.attachments) == 2 + assert instance.attachments[0] is pointA + assert instance.attachments[1] is pointB + assert isinstance(instance.attachments[0], Point) + assert instance.attachments[0].name == 'pA' + assert isinstance(instance.attachments[1], Point) + assert instance.attachments[1].name == 'pB' + + @staticmethod + @pytest.mark.parametrize( + 'attachments', + [ + (Point('pA'), ), + (Point('pA'), Point('pB'), Point('pZ')), + ] + ) + def test_invalid_attachments_incorrect_number(attachments): + with pytest.raises(ValueError): + _ = LinearPathway(*attachments) + + @staticmethod + @pytest.mark.parametrize( + 'attachments', + [ + (None, Point('pB')), + (Point('pA'), None), + ] + ) + def test_invalid_attachments_not_point(attachments): + with pytest.raises(TypeError): + _ = LinearPathway(*attachments) + + @pytest.fixture(autouse=True) + def _linear_pathway_fixture(self): + self.N = ReferenceFrame('N') + self.pA = Point('pA') + self.pB = Point('pB') + self.pathway = LinearPathway(self.pA, self.pB) + self.q1 = dynamicsymbols('q1') + self.q2 = dynamicsymbols('q2') + self.q3 = dynamicsymbols('q3') + self.q1d = dynamicsymbols('q1', 1) + self.q2d = dynamicsymbols('q2', 1) + self.q3d = dynamicsymbols('q3', 1) + self.F = Symbol('F') + + def test_properties_are_immutable(self): + instance = LinearPathway(self.pA, self.pB) + with pytest.raises(AttributeError): + instance.attachments = None + with pytest.raises(TypeError): + instance.attachments[0] = None + with pytest.raises(TypeError): + instance.attachments[1] = None + + def test_repr(self): + pathway = LinearPathway(self.pA, self.pB) + expected = 'LinearPathway(pA, pB)' + assert repr(pathway) == expected + + def test_static_pathway_length(self): + self.pB.set_pos(self.pA, 2*self.N.x) + assert self.pathway.length == 2 + + def test_static_pathway_extension_velocity(self): + self.pB.set_pos(self.pA, 2*self.N.x) + assert self.pathway.extension_velocity == 0 + + def test_static_pathway_to_loads(self): + self.pB.set_pos(self.pA, 2*self.N.x) + expected = [ + (self.pA, - self.F*self.N.x), + (self.pB, self.F*self.N.x), + ] + assert self.pathway.to_loads(self.F) == expected + + def test_2D_pathway_length(self): + self.pB.set_pos(self.pA, 2*self.q1*self.N.x) + expected = 2*sqrt(self.q1**2) + assert self.pathway.length == expected + + def test_2D_pathway_extension_velocity(self): + self.pB.set_pos(self.pA, 2*self.q1*self.N.x) + expected = 2*sqrt(self.q1**2)*self.q1d/self.q1 + assert self.pathway.extension_velocity == expected + + def test_2D_pathway_to_loads(self): + self.pB.set_pos(self.pA, 2*self.q1*self.N.x) + expected = [ + (self.pA, - self.F*(self.q1 / sqrt(self.q1**2))*self.N.x), + (self.pB, self.F*(self.q1 / sqrt(self.q1**2))*self.N.x), + ] + assert self.pathway.to_loads(self.F) == expected + + def test_3D_pathway_length(self): + self.pB.set_pos( + self.pA, + self.q1*self.N.x - self.q2*self.N.y + 2*self.q3*self.N.z, + ) + expected = sqrt(self.q1**2 + self.q2**2 + 4*self.q3**2) + assert simplify(self.pathway.length - expected) == 0 + + def test_3D_pathway_extension_velocity(self): + self.pB.set_pos( + self.pA, + self.q1*self.N.x - self.q2*self.N.y + 2*self.q3*self.N.z, + ) + length = sqrt(self.q1**2 + self.q2**2 + 4*self.q3**2) + expected = ( + self.q1*self.q1d/length + + self.q2*self.q2d/length + + 4*self.q3*self.q3d/length + ) + assert simplify(self.pathway.extension_velocity - expected) == 0 + + def test_3D_pathway_to_loads(self): + self.pB.set_pos( + self.pA, + self.q1*self.N.x - self.q2*self.N.y + 2*self.q3*self.N.z, + ) + length = sqrt(self.q1**2 + self.q2**2 + 4*self.q3**2) + pO_force = ( + - self.F*self.q1*self.N.x/length + + self.F*self.q2*self.N.y/length + - 2*self.F*self.q3*self.N.z/length + ) + pI_force = ( + self.F*self.q1*self.N.x/length + - self.F*self.q2*self.N.y/length + + 2*self.F*self.q3*self.N.z/length + ) + expected = [ + (self.pA, pO_force), + (self.pB, pI_force), + ] + assert self.pathway.to_loads(self.F) == expected + + +class TestObstacleSetPathway: + + def test_is_pathway_base_subclass(self): + assert issubclass(ObstacleSetPathway, PathwayBase) + + @staticmethod + @pytest.mark.parametrize( + 'num_attachments, attachments', + [ + (3, [Point(name) for name in ('pO', 'pA', 'pI')]), + (4, [Point(name) for name in ('pO', 'pA', 'pB', 'pI')]), + (5, [Point(name) for name in ('pO', 'pA', 'pB', 'pC', 'pI')]), + (6, [Point(name) for name in ('pO', 'pA', 'pB', 'pC', 'pD', 'pI')]), + ] + ) + def test_valid_constructor(num_attachments, attachments): + instance = ObstacleSetPathway(*attachments) + assert isinstance(instance, ObstacleSetPathway) + assert hasattr(instance, 'attachments') + assert len(instance.attachments) == num_attachments + for attachment in instance.attachments: + assert isinstance(attachment, Point) + + @staticmethod + @pytest.mark.parametrize( + 'attachments', + [[Point('pO')], [Point('pO'), Point('pI')]], + ) + def test_invalid_constructor_attachments_incorrect_number(attachments): + with pytest.raises(ValueError): + _ = ObstacleSetPathway(*attachments) + + @staticmethod + @pytest.mark.parametrize( + 'attachments', + [ + (None, Point('pA'), Point('pI')), + (Point('pO'), None, Point('pI')), + (Point('pO'), Point('pA'), None), + ] + ) + def test_invalid_constructor_attachments_not_point(attachments): + with pytest.raises(TypeError): + _ = WrappingPathway(*attachments) # type: ignore + + def test_properties_are_immutable(self): + pathway = ObstacleSetPathway(Point('pO'), Point('pA'), Point('pI')) + with pytest.raises(AttributeError): + pathway.attachments = None # type: ignore + with pytest.raises(TypeError): + pathway.attachments[0] = None # type: ignore + with pytest.raises(TypeError): + pathway.attachments[1] = None # type: ignore + with pytest.raises(TypeError): + pathway.attachments[-1] = None # type: ignore + + @staticmethod + @pytest.mark.parametrize( + 'attachments, expected', + [ + ( + [Point(name) for name in ('pO', 'pA', 'pI')], + 'ObstacleSetPathway(pO, pA, pI)' + ), + ( + [Point(name) for name in ('pO', 'pA', 'pB', 'pI')], + 'ObstacleSetPathway(pO, pA, pB, pI)' + ), + ( + [Point(name) for name in ('pO', 'pA', 'pB', 'pC', 'pI')], + 'ObstacleSetPathway(pO, pA, pB, pC, pI)' + ), + ] + ) + def test_repr(attachments, expected): + pathway = ObstacleSetPathway(*attachments) + assert repr(pathway) == expected + + @pytest.fixture(autouse=True) + def _obstacle_set_pathway_fixture(self): + self.N = ReferenceFrame('N') + self.pO = Point('pO') + self.pI = Point('pI') + self.pA = Point('pA') + self.pB = Point('pB') + self.q = dynamicsymbols('q') + self.qd = dynamicsymbols('q', 1) + self.F = Symbol('F') + + def test_static_pathway_length(self): + self.pA.set_pos(self.pO, self.N.x) + self.pB.set_pos(self.pO, self.N.y) + self.pI.set_pos(self.pO, self.N.z) + pathway = ObstacleSetPathway(self.pO, self.pA, self.pB, self.pI) + assert pathway.length == 1 + 2 * sqrt(2) + + def test_static_pathway_extension_velocity(self): + self.pA.set_pos(self.pO, self.N.x) + self.pB.set_pos(self.pO, self.N.y) + self.pI.set_pos(self.pO, self.N.z) + pathway = ObstacleSetPathway(self.pO, self.pA, self.pB, self.pI) + assert pathway.extension_velocity == 0 + + def test_static_pathway_to_loads(self): + self.pA.set_pos(self.pO, self.N.x) + self.pB.set_pos(self.pO, self.N.y) + self.pI.set_pos(self.pO, self.N.z) + pathway = ObstacleSetPathway(self.pO, self.pA, self.pB, self.pI) + expected = [ + Force(self.pO, -self.F * self.N.x), + Force(self.pA, self.F * self.N.x), + Force(self.pA, self.F * sqrt(2) / 2 * (self.N.x - self.N.y)), + Force(self.pB, self.F * sqrt(2) / 2 * (self.N.y - self.N.x)), + Force(self.pB, self.F * sqrt(2) / 2 * (self.N.y - self.N.z)), + Force(self.pI, self.F * sqrt(2) / 2 * (self.N.z - self.N.y)), + ] + assert pathway.to_loads(self.F) == expected + + def test_2D_pathway_length(self): + self.pA.set_pos(self.pO, -(self.N.x + self.N.y)) + self.pB.set_pos( + self.pO, cos(self.q) * self.N.x - (sin(self.q) + 1) * self.N.y + ) + self.pI.set_pos( + self.pO, sin(self.q) * self.N.x + (cos(self.q) - 1) * self.N.y + ) + pathway = ObstacleSetPathway(self.pO, self.pA, self.pB, self.pI) + expected = 2 * sqrt(2) + sqrt(2 + 2*cos(self.q)) + assert (pathway.length - expected).simplify() == 0 + + def test_2D_pathway_extension_velocity(self): + self.pA.set_pos(self.pO, -(self.N.x + self.N.y)) + self.pB.set_pos( + self.pO, cos(self.q) * self.N.x - (sin(self.q) + 1) * self.N.y + ) + self.pI.set_pos( + self.pO, sin(self.q) * self.N.x + (cos(self.q) - 1) * self.N.y + ) + pathway = ObstacleSetPathway(self.pO, self.pA, self.pB, self.pI) + expected = - (sqrt(2) * sin(self.q) * self.qd) / (2 * sqrt(cos(self.q) + 1)) + assert (pathway.extension_velocity - expected).simplify() == 0 + + def test_2D_pathway_to_loads(self): + self.pA.set_pos(self.pO, -(self.N.x + self.N.y)) + self.pB.set_pos( + self.pO, cos(self.q) * self.N.x - (sin(self.q) + 1) * self.N.y + ) + self.pI.set_pos( + self.pO, sin(self.q) * self.N.x + (cos(self.q) - 1) * self.N.y + ) + pathway = ObstacleSetPathway(self.pO, self.pA, self.pB, self.pI) + pO_pA_force_vec = sqrt(2) / 2 * (self.N.x + self.N.y) + pA_pB_force_vec = ( + - sqrt(2 * cos(self.q) + 2) / 2 * self.N.x + + sqrt(2) * sin(self.q) / (2 * sqrt(cos(self.q) + 1)) * self.N.y + ) + pB_pI_force_vec = cos(self.q + pi/4) * self.N.x - sin(self.q + pi/4) * self.N.y + expected = [ + Force(self.pO, self.F * pO_pA_force_vec), + Force(self.pA, -self.F * pO_pA_force_vec), + Force(self.pA, self.F * pA_pB_force_vec), + Force(self.pB, -self.F * pA_pB_force_vec), + Force(self.pB, self.F * pB_pI_force_vec), + Force(self.pI, -self.F * pB_pI_force_vec), + ] + assert _simplify_loads(pathway.to_loads(self.F)) == expected + + +class TestWrappingPathway: + + def test_is_pathway_base_subclass(self): + assert issubclass(WrappingPathway, PathwayBase) + + @pytest.fixture(autouse=True) + def _wrapping_pathway_fixture(self): + self.pA = Point('pA') + self.pB = Point('pB') + self.r = Symbol('r', positive=True) + self.pO = Point('pO') + self.N = ReferenceFrame('N') + self.ax = self.N.z + self.sphere = WrappingSphere(self.r, self.pO) + self.cylinder = WrappingCylinder(self.r, self.pO, self.ax) + self.pathway = WrappingPathway(self.pA, self.pB, self.cylinder) + self.F = Symbol('F') + + def test_valid_constructor(self): + instance = WrappingPathway(self.pA, self.pB, self.cylinder) + assert isinstance(instance, WrappingPathway) + assert hasattr(instance, 'attachments') + assert len(instance.attachments) == 2 + assert isinstance(instance.attachments[0], Point) + assert instance.attachments[0] == self.pA + assert isinstance(instance.attachments[1], Point) + assert instance.attachments[1] == self.pB + assert hasattr(instance, 'geometry') + assert isinstance(instance.geometry, WrappingGeometryBase) + assert instance.geometry == self.cylinder + + @pytest.mark.parametrize( + 'attachments', + [ + (Point('pA'), ), + (Point('pA'), Point('pB'), Point('pZ')), + ] + ) + def test_invalid_constructor_attachments_incorrect_number(self, attachments): + with pytest.raises(TypeError): + _ = WrappingPathway(*attachments, self.cylinder) + + @staticmethod + @pytest.mark.parametrize( + 'attachments', + [ + (None, Point('pB')), + (Point('pA'), None), + ] + ) + def test_invalid_constructor_attachments_not_point(attachments): + with pytest.raises(TypeError): + _ = WrappingPathway(*attachments) + + def test_invalid_constructor_geometry_is_not_supplied(self): + with pytest.raises(TypeError): + _ = WrappingPathway(self.pA, self.pB) + + @pytest.mark.parametrize( + 'geometry', + [ + Symbol('r'), + dynamicsymbols('q'), + ReferenceFrame('N'), + ReferenceFrame('N').x, + ] + ) + def test_invalid_geometry_not_geometry(self, geometry): + with pytest.raises(TypeError): + _ = WrappingPathway(self.pA, self.pB, geometry) + + def test_attachments_property_is_immutable(self): + with pytest.raises(TypeError): + self.pathway.attachments[0] = self.pB + with pytest.raises(TypeError): + self.pathway.attachments[1] = self.pA + + def test_geometry_property_is_immutable(self): + with pytest.raises(AttributeError): + self.pathway.geometry = None + + def test_repr(self): + expected = ( + f'WrappingPathway(pA, pB, ' + f'geometry={self.cylinder!r})' + ) + assert repr(self.pathway) == expected + + @staticmethod + def _expand_pos_to_vec(pos, frame): + return sum(mag*unit for (mag, unit) in zip(pos, frame)) + + @pytest.mark.parametrize( + 'pA_vec, pB_vec, factor', + [ + ((1, 0, 0), (0, 1, 0), pi/2), + ((0, 1, 0), (sqrt(2)/2, -sqrt(2)/2, 0), 3*pi/4), + ((1, 0, 0), (Rational(1, 2), sqrt(3)/2, 0), pi/3), + ] + ) + def test_static_pathway_on_sphere_length(self, pA_vec, pB_vec, factor): + pA_vec = self._expand_pos_to_vec(pA_vec, self.N) + pB_vec = self._expand_pos_to_vec(pB_vec, self.N) + self.pA.set_pos(self.pO, self.r*pA_vec) + self.pB.set_pos(self.pO, self.r*pB_vec) + pathway = WrappingPathway(self.pA, self.pB, self.sphere) + expected = factor*self.r + assert simplify(pathway.length - expected) == 0 + + @pytest.mark.parametrize( + 'pA_vec, pB_vec, factor', + [ + ((1, 0, 0), (0, 1, 0), Rational(1, 2)*pi), + ((1, 0, 0), (-1, 0, 0), pi), + ((-1, 0, 0), (1, 0, 0), pi), + ((0, 1, 0), (sqrt(2)/2, -sqrt(2)/2, 0), 5*pi/4), + ((1, 0, 0), (Rational(1, 2), sqrt(3)/2, 0), pi/3), + ( + (0, 1, 0), + (sqrt(2)*Rational(1, 2), -sqrt(2)*Rational(1, 2), 1), + sqrt(1 + (Rational(5, 4)*pi)**2), + ), + ( + (1, 0, 0), + (Rational(1, 2), sqrt(3)*Rational(1, 2), 1), + sqrt(1 + (Rational(1, 3)*pi)**2), + ), + ] + ) + def test_static_pathway_on_cylinder_length(self, pA_vec, pB_vec, factor): + pA_vec = self._expand_pos_to_vec(pA_vec, self.N) + pB_vec = self._expand_pos_to_vec(pB_vec, self.N) + self.pA.set_pos(self.pO, self.r*pA_vec) + self.pB.set_pos(self.pO, self.r*pB_vec) + pathway = WrappingPathway(self.pA, self.pB, self.cylinder) + expected = factor*sqrt(self.r**2) + assert simplify(pathway.length - expected) == 0 + + @pytest.mark.parametrize( + 'pA_vec, pB_vec', + [ + ((1, 0, 0), (0, 1, 0)), + ((0, 1, 0), (sqrt(2)*Rational(1, 2), -sqrt(2)*Rational(1, 2), 0)), + ((1, 0, 0), (Rational(1, 2), sqrt(3)*Rational(1, 2), 0)), + ] + ) + def test_static_pathway_on_sphere_extension_velocity(self, pA_vec, pB_vec): + pA_vec = self._expand_pos_to_vec(pA_vec, self.N) + pB_vec = self._expand_pos_to_vec(pB_vec, self.N) + self.pA.set_pos(self.pO, self.r*pA_vec) + self.pB.set_pos(self.pO, self.r*pB_vec) + pathway = WrappingPathway(self.pA, self.pB, self.sphere) + assert pathway.extension_velocity == 0 + + @pytest.mark.parametrize( + 'pA_vec, pB_vec', + [ + ((1, 0, 0), (0, 1, 0)), + ((1, 0, 0), (-1, 0, 0)), + ((-1, 0, 0), (1, 0, 0)), + ((0, 1, 0), (sqrt(2)/2, -sqrt(2)/2, 0)), + ((1, 0, 0), (Rational(1, 2), sqrt(3)/2, 0)), + ((0, 1, 0), (sqrt(2)*Rational(1, 2), -sqrt(2)/2, 1)), + ((1, 0, 0), (Rational(1, 2), sqrt(3)/2, 1)), + ] + ) + def test_static_pathway_on_cylinder_extension_velocity(self, pA_vec, pB_vec): + pA_vec = self._expand_pos_to_vec(pA_vec, self.N) + pB_vec = self._expand_pos_to_vec(pB_vec, self.N) + self.pA.set_pos(self.pO, self.r*pA_vec) + self.pB.set_pos(self.pO, self.r*pB_vec) + pathway = WrappingPathway(self.pA, self.pB, self.cylinder) + assert pathway.extension_velocity == 0 + + @pytest.mark.parametrize( + 'pA_vec, pB_vec, pA_vec_expected, pB_vec_expected, pO_vec_expected', + ( + ((1, 0, 0), (0, 1, 0), (0, 1, 0), (1, 0, 0), (-1, -1, 0)), + ( + (0, 1, 0), + (sqrt(2)/2, -sqrt(2)/2, 0), + (1, 0, 0), + (sqrt(2)/2, sqrt(2)/2, 0), + (-1 - sqrt(2)/2, -sqrt(2)/2, 0) + ), + ( + (1, 0, 0), + (Rational(1, 2), sqrt(3)/2, 0), + (0, 1, 0), + (sqrt(3)/2, -Rational(1, 2), 0), + (-sqrt(3)/2, Rational(1, 2) - 1, 0), + ), + ) + ) + def test_static_pathway_on_sphere_to_loads( + self, + pA_vec, + pB_vec, + pA_vec_expected, + pB_vec_expected, + pO_vec_expected, + ): + pA_vec = self._expand_pos_to_vec(pA_vec, self.N) + pB_vec = self._expand_pos_to_vec(pB_vec, self.N) + self.pA.set_pos(self.pO, self.r*pA_vec) + self.pB.set_pos(self.pO, self.r*pB_vec) + pathway = WrappingPathway(self.pA, self.pB, self.sphere) + + pA_vec_expected = sum( + mag*unit for (mag, unit) in zip(pA_vec_expected, self.N) + ) + pB_vec_expected = sum( + mag*unit for (mag, unit) in zip(pB_vec_expected, self.N) + ) + pO_vec_expected = sum( + mag*unit for (mag, unit) in zip(pO_vec_expected, self.N) + ) + expected = [ + Force(self.pA, self.F*(self.r**3/sqrt(self.r**6))*pA_vec_expected), + Force(self.pB, self.F*(self.r**3/sqrt(self.r**6))*pB_vec_expected), + Force(self.pO, self.F*(self.r**3/sqrt(self.r**6))*pO_vec_expected), + ] + assert pathway.to_loads(self.F) == expected + + @pytest.mark.parametrize( + 'pA_vec, pB_vec, pA_vec_expected, pB_vec_expected, pO_vec_expected', + ( + ((1, 0, 0), (0, 1, 0), (0, 1, 0), (1, 0, 0), (-1, -1, 0)), + ((1, 0, 0), (-1, 0, 0), (0, 1, 0), (0, 1, 0), (0, -2, 0)), + ((-1, 0, 0), (1, 0, 0), (0, -1, 0), (0, -1, 0), (0, 2, 0)), + ( + (0, 1, 0), + (sqrt(2)/2, -sqrt(2)/2, 0), + (-1, 0, 0), + (-sqrt(2)/2, -sqrt(2)/2, 0), + (1 + sqrt(2)/2, sqrt(2)/2, 0) + ), + ( + (1, 0, 0), + (Rational(1, 2), sqrt(3)/2, 0), + (0, 1, 0), + (sqrt(3)/2, -Rational(1, 2), 0), + (-sqrt(3)/2, Rational(1, 2) - 1, 0), + ), + ( + (1, 0, 0), + (sqrt(2)/2, sqrt(2)/2, 0), + (0, 1, 0), + (sqrt(2)/2, -sqrt(2)/2, 0), + (-sqrt(2)/2, sqrt(2)/2 - 1, 0), + ), + ((0, 1, 0), (0, 1, 1), (0, 0, 1), (0, 0, -1), (0, 0, 0)), + ( + (0, 1, 0), + (sqrt(2)/2, -sqrt(2)/2, 1), + (-5*pi/sqrt(16 + 25*pi**2), 0, 4/sqrt(16 + 25*pi**2)), + ( + -5*sqrt(2)*pi/(2*sqrt(16 + 25*pi**2)), + -5*sqrt(2)*pi/(2*sqrt(16 + 25*pi**2)), + -4/sqrt(16 + 25*pi**2), + ), + ( + 5*(sqrt(2) + 2)*pi/(2*sqrt(16 + 25*pi**2)), + 5*sqrt(2)*pi/(2*sqrt(16 + 25*pi**2)), + 0, + ), + ), + ) + ) + def test_static_pathway_on_cylinder_to_loads( + self, + pA_vec, + pB_vec, + pA_vec_expected, + pB_vec_expected, + pO_vec_expected, + ): + pA_vec = self._expand_pos_to_vec(pA_vec, self.N) + pB_vec = self._expand_pos_to_vec(pB_vec, self.N) + self.pA.set_pos(self.pO, self.r*pA_vec) + self.pB.set_pos(self.pO, self.r*pB_vec) + pathway = WrappingPathway(self.pA, self.pB, self.cylinder) + + pA_force_expected = self.F*self._expand_pos_to_vec(pA_vec_expected, + self.N) + pB_force_expected = self.F*self._expand_pos_to_vec(pB_vec_expected, + self.N) + pO_force_expected = self.F*self._expand_pos_to_vec(pO_vec_expected, + self.N) + expected = [ + Force(self.pA, pA_force_expected), + Force(self.pB, pB_force_expected), + Force(self.pO, pO_force_expected), + ] + assert _simplify_loads(pathway.to_loads(self.F)) == expected + + def test_2D_pathway_on_cylinder_length(self): + q = dynamicsymbols('q') + pA_pos = self.r*self.N.x + pB_pos = self.r*(cos(q)*self.N.x + sin(q)*self.N.y) + self.pA.set_pos(self.pO, pA_pos) + self.pB.set_pos(self.pO, pB_pos) + expected = self.r*sqrt(q**2) + assert simplify(self.pathway.length - expected) == 0 + + def test_2D_pathway_on_cylinder_extension_velocity(self): + q = dynamicsymbols('q') + qd = dynamicsymbols('q', 1) + pA_pos = self.r*self.N.x + pB_pos = self.r*(cos(q)*self.N.x + sin(q)*self.N.y) + self.pA.set_pos(self.pO, pA_pos) + self.pB.set_pos(self.pO, pB_pos) + expected = self.r*(sqrt(q**2)/q)*qd + assert simplify(self.pathway.extension_velocity - expected) == 0 + + def test_2D_pathway_on_cylinder_to_loads(self): + q = dynamicsymbols('q') + pA_pos = self.r*self.N.x + pB_pos = self.r*(cos(q)*self.N.x + sin(q)*self.N.y) + self.pA.set_pos(self.pO, pA_pos) + self.pB.set_pos(self.pO, pB_pos) + + pA_force = self.F*self.N.y + pB_force = self.F*(sin(q)*self.N.x - cos(q)*self.N.y) + pO_force = self.F*(-sin(q)*self.N.x + (cos(q) - 1)*self.N.y) + expected = [ + Force(self.pA, pA_force), + Force(self.pB, pB_force), + Force(self.pO, pO_force), + ] + + loads = _simplify_loads(self.pathway.to_loads(self.F)) + assert loads == expected diff --git a/lib/python3.12/site-packages/sympy/physics/mechanics/tests/test_rigidbody.py b/lib/python3.12/site-packages/sympy/physics/mechanics/tests/test_rigidbody.py new file mode 100644 index 0000000000000000000000000000000000000000..78161e0c9fc33be6e3d274034b67278c8ceee8fd --- /dev/null +++ b/lib/python3.12/site-packages/sympy/physics/mechanics/tests/test_rigidbody.py @@ -0,0 +1,184 @@ +from sympy.physics.mechanics import Point, ReferenceFrame, Dyadic, RigidBody +from sympy.physics.mechanics import dynamicsymbols, outer, inertia, Inertia +from sympy.physics.mechanics import inertia_of_point_mass +from sympy import expand, zeros, simplify, symbols +from sympy.testing.pytest import raises, warns_deprecated_sympy + + +def test_rigidbody_default(): + # Test default + b = RigidBody('B') + I = inertia(b.frame, *symbols('B_ixx B_iyy B_izz B_ixy B_iyz B_izx')) + assert b.name == 'B' + assert b.mass == symbols('B_mass') + assert b.masscenter.name == 'B_masscenter' + assert b.inertia == (I, b.masscenter) + assert b.central_inertia == I + assert b.frame.name == 'B_frame' + assert b.__str__() == 'B' + assert b.__repr__() == ( + "RigidBody('B', masscenter=B_masscenter, frame=B_frame, mass=B_mass, " + "inertia=Inertia(dyadic=B_ixx*(B_frame.x|B_frame.x) + " + "B_ixy*(B_frame.x|B_frame.y) + B_izx*(B_frame.x|B_frame.z) + " + "B_ixy*(B_frame.y|B_frame.x) + B_iyy*(B_frame.y|B_frame.y) + " + "B_iyz*(B_frame.y|B_frame.z) + B_izx*(B_frame.z|B_frame.x) + " + "B_iyz*(B_frame.z|B_frame.y) + B_izz*(B_frame.z|B_frame.z), " + "point=B_masscenter))") + + +def test_rigidbody(): + m, m2, v1, v2, v3, omega = symbols('m m2 v1 v2 v3 omega') + A = ReferenceFrame('A') + A2 = ReferenceFrame('A2') + P = Point('P') + P2 = Point('P2') + I = Dyadic(0) + I2 = Dyadic(0) + B = RigidBody('B', P, A, m, (I, P)) + assert B.mass == m + assert B.frame == A + assert B.masscenter == P + assert B.inertia == (I, B.masscenter) + + B.mass = m2 + B.frame = A2 + B.masscenter = P2 + B.inertia = (I2, B.masscenter) + raises(TypeError, lambda: RigidBody(P, P, A, m, (I, P))) + raises(TypeError, lambda: RigidBody('B', P, P, m, (I, P))) + raises(TypeError, lambda: RigidBody('B', P, A, m, (P, P))) + raises(TypeError, lambda: RigidBody('B', P, A, m, (I, I))) + assert B.__str__() == 'B' + assert B.mass == m2 + assert B.frame == A2 + assert B.masscenter == P2 + assert B.inertia == (I2, B.masscenter) + assert isinstance(B.inertia, Inertia) + + # Testing linear momentum function assuming A2 is the inertial frame + N = ReferenceFrame('N') + P2.set_vel(N, v1 * N.x + v2 * N.y + v3 * N.z) + assert B.linear_momentum(N) == m2 * (v1 * N.x + v2 * N.y + v3 * N.z) + + +def test_rigidbody2(): + M, v, r, omega, g, h = dynamicsymbols('M v r omega g h') + N = ReferenceFrame('N') + b = ReferenceFrame('b') + b.set_ang_vel(N, omega * b.x) + P = Point('P') + I = outer(b.x, b.x) + Inertia_tuple = (I, P) + B = RigidBody('B', P, b, M, Inertia_tuple) + P.set_vel(N, v * b.x) + assert B.angular_momentum(P, N) == omega * b.x + O = Point('O') + O.set_vel(N, v * b.x) + P.set_pos(O, r * b.y) + assert B.angular_momentum(O, N) == omega * b.x - M*v*r*b.z + B.potential_energy = M * g * h + assert B.potential_energy == M * g * h + assert expand(2 * B.kinetic_energy(N)) == omega**2 + M * v**2 + + +def test_rigidbody3(): + q1, q2, q3, q4 = dynamicsymbols('q1:5') + p1, p2, p3 = symbols('p1:4') + m = symbols('m') + + A = ReferenceFrame('A') + B = A.orientnew('B', 'axis', [q1, A.x]) + O = Point('O') + O.set_vel(A, q2*A.x + q3*A.y + q4*A.z) + P = O.locatenew('P', p1*B.x + p2*B.y + p3*B.z) + P.v2pt_theory(O, A, B) + I = outer(B.x, B.x) + + rb1 = RigidBody('rb1', P, B, m, (I, P)) + # I_S/O = I_S/S* + I_S*/O + rb2 = RigidBody('rb2', P, B, m, + (I + inertia_of_point_mass(m, P.pos_from(O), B), O)) + + assert rb1.central_inertia == rb2.central_inertia + assert rb1.angular_momentum(O, A) == rb2.angular_momentum(O, A) + + +def test_pendulum_angular_momentum(): + """Consider a pendulum of length OA = 2a, of mass m as a rigid body of + center of mass G (OG = a) which turn around (O,z). The angle between the + reference frame R and the rod is q. The inertia of the body is I = + (G,0,ma^2/3,ma^2/3). """ + + m, a = symbols('m, a') + q = dynamicsymbols('q') + + R = ReferenceFrame('R') + R1 = R.orientnew('R1', 'Axis', [q, R.z]) + R1.set_ang_vel(R, q.diff() * R.z) + + I = inertia(R1, 0, m * a**2 / 3, m * a**2 / 3) + + O = Point('O') + + A = O.locatenew('A', 2*a * R1.x) + G = O.locatenew('G', a * R1.x) + + S = RigidBody('S', G, R1, m, (I, G)) + + O.set_vel(R, 0) + A.v2pt_theory(O, R, R1) + G.v2pt_theory(O, R, R1) + + assert (4 * m * a**2 / 3 * q.diff() * R.z - + S.angular_momentum(O, R).express(R)) == 0 + + +def test_rigidbody_inertia(): + N = ReferenceFrame('N') + m, Ix, Iy, Iz, a, b = symbols('m, I_x, I_y, I_z, a, b') + Io = inertia(N, Ix, Iy, Iz) + o = Point('o') + p = o.locatenew('p', a * N.x + b * N.y) + R = RigidBody('R', o, N, m, (Io, p)) + I_check = inertia(N, Ix - b ** 2 * m, Iy - a ** 2 * m, + Iz - m * (a ** 2 + b ** 2), m * a * b) + assert isinstance(R.inertia, Inertia) + assert R.inertia == (Io, p) + assert R.central_inertia == I_check + R.central_inertia = Io + assert R.inertia == (Io, o) + assert R.central_inertia == Io + R.inertia = (Io, p) + assert R.inertia == (Io, p) + assert R.central_inertia == I_check + # parse Inertia object + R.inertia = Inertia(Io, o) + assert R.inertia == (Io, o) + + +def test_parallel_axis(): + N = ReferenceFrame('N') + m, Ix, Iy, Iz, a, b = symbols('m, I_x, I_y, I_z, a, b') + Io = inertia(N, Ix, Iy, Iz) + o = Point('o') + p = o.locatenew('p', a * N.x + b * N.y) + R = RigidBody('R', o, N, m, (Io, o)) + Ip = R.parallel_axis(p) + Ip_expected = inertia(N, Ix + m * b**2, Iy + m * a**2, + Iz + m * (a**2 + b**2), ixy=-m * a * b) + assert Ip == Ip_expected + # Reference frame from which the parallel axis is viewed should not matter + A = ReferenceFrame('A') + A.orient_axis(N, N.z, 1) + assert simplify( + (R.parallel_axis(p, A) - Ip_expected).to_matrix(A)) == zeros(3, 3) + + +def test_deprecated_set_potential_energy(): + m, g, h = symbols('m g h') + A = ReferenceFrame('A') + P = Point('P') + I = Dyadic(0) + B = RigidBody('B', P, A, m, (I, P)) + with warns_deprecated_sympy(): + B.set_potential_energy(m*g*h) diff --git a/lib/python3.12/site-packages/sympy/physics/mechanics/tests/test_system.py b/lib/python3.12/site-packages/sympy/physics/mechanics/tests/test_system.py new file mode 100644 index 0000000000000000000000000000000000000000..6fdac1ea10e9f71f8cf999cc5069da7567f67adf --- /dev/null +++ b/lib/python3.12/site-packages/sympy/physics/mechanics/tests/test_system.py @@ -0,0 +1,245 @@ +from sympy import symbols, Matrix, atan, zeros +from sympy.simplify.simplify import simplify +from sympy.physics.mechanics import (dynamicsymbols, Particle, Point, + ReferenceFrame, SymbolicSystem) +from sympy.testing.pytest import raises + +# This class is going to be tested using a simple pendulum set up in x and y +# coordinates +x, y, u, v, lam = dynamicsymbols('x y u v lambda') +m, l, g = symbols('m l g') + +# Set up the different forms the equations can take +# [1] Explicit form where the kinematics and dynamics are combined +# x' = F(x, t, r, p) +# +# [2] Implicit form where the kinematics and dynamics are combined +# M(x, p) x' = F(x, t, r, p) +# +# [3] Implicit form where the kinematics and dynamics are separate +# M(q, p) u' = F(q, u, t, r, p) +# q' = G(q, u, t, r, p) +dyn_implicit_mat = Matrix([[1, 0, -x/m], + [0, 1, -y/m], + [0, 0, l**2/m]]) + +dyn_implicit_rhs = Matrix([0, 0, u**2 + v**2 - g*y]) + +comb_implicit_mat = Matrix([[1, 0, 0, 0, 0], + [0, 1, 0, 0, 0], + [0, 0, 1, 0, -x/m], + [0, 0, 0, 1, -y/m], + [0, 0, 0, 0, l**2/m]]) + +comb_implicit_rhs = Matrix([u, v, 0, 0, u**2 + v**2 - g*y]) + +kin_explicit_rhs = Matrix([u, v]) + +comb_explicit_rhs = comb_implicit_mat.LUsolve(comb_implicit_rhs) + +# Set up a body and load to pass into the system +theta = atan(x/y) +N = ReferenceFrame('N') +A = N.orientnew('A', 'Axis', [theta, N.z]) +O = Point('O') +P = O.locatenew('P', l * A.x) + +Pa = Particle('Pa', P, m) + +bodies = [Pa] +loads = [(P, g * m * N.x)] + +# Set up some output equations to be given to SymbolicSystem +# Change to make these fit the pendulum +PE = symbols("PE") +out_eqns = {PE: m*g*(l+y)} + +# Set up remaining arguments that can be passed to SymbolicSystem +alg_con = [2] +alg_con_full = [4] +coordinates = (x, y, lam) +speeds = (u, v) +states = (x, y, u, v, lam) +coord_idxs = (0, 1) +speed_idxs = (2, 3) + + +def test_form_1(): + symsystem1 = SymbolicSystem(states, comb_explicit_rhs, + alg_con=alg_con_full, output_eqns=out_eqns, + coord_idxs=coord_idxs, speed_idxs=speed_idxs, + bodies=bodies, loads=loads) + + assert symsystem1.coordinates == Matrix([x, y]) + assert symsystem1.speeds == Matrix([u, v]) + assert symsystem1.states == Matrix([x, y, u, v, lam]) + + assert symsystem1.alg_con == [4] + + inter = comb_explicit_rhs + assert simplify(symsystem1.comb_explicit_rhs - inter) == zeros(5, 1) + + assert set(symsystem1.dynamic_symbols()) == {y, v, lam, u, x} + assert type(symsystem1.dynamic_symbols()) == tuple + assert set(symsystem1.constant_symbols()) == {l, g, m} + assert type(symsystem1.constant_symbols()) == tuple + + assert symsystem1.output_eqns == out_eqns + + assert symsystem1.bodies == (Pa,) + assert symsystem1.loads == ((P, g * m * N.x),) + + +def test_form_2(): + symsystem2 = SymbolicSystem(coordinates, comb_implicit_rhs, speeds=speeds, + mass_matrix=comb_implicit_mat, + alg_con=alg_con_full, output_eqns=out_eqns, + bodies=bodies, loads=loads) + + assert symsystem2.coordinates == Matrix([x, y, lam]) + assert symsystem2.speeds == Matrix([u, v]) + assert symsystem2.states == Matrix([x, y, lam, u, v]) + + assert symsystem2.alg_con == [4] + + inter = comb_implicit_rhs + assert simplify(symsystem2.comb_implicit_rhs - inter) == zeros(5, 1) + assert simplify(symsystem2.comb_implicit_mat-comb_implicit_mat) == zeros(5) + + assert set(symsystem2.dynamic_symbols()) == {y, v, lam, u, x} + assert type(symsystem2.dynamic_symbols()) == tuple + assert set(symsystem2.constant_symbols()) == {l, g, m} + assert type(symsystem2.constant_symbols()) == tuple + + inter = comb_explicit_rhs + symsystem2.compute_explicit_form() + assert simplify(symsystem2.comb_explicit_rhs - inter) == zeros(5, 1) + + + assert symsystem2.output_eqns == out_eqns + + assert symsystem2.bodies == (Pa,) + assert symsystem2.loads == ((P, g * m * N.x),) + + +def test_form_3(): + symsystem3 = SymbolicSystem(states, dyn_implicit_rhs, + mass_matrix=dyn_implicit_mat, + coordinate_derivatives=kin_explicit_rhs, + alg_con=alg_con, coord_idxs=coord_idxs, + speed_idxs=speed_idxs, bodies=bodies, + loads=loads) + + assert symsystem3.coordinates == Matrix([x, y]) + assert symsystem3.speeds == Matrix([u, v]) + assert symsystem3.states == Matrix([x, y, u, v, lam]) + + assert symsystem3.alg_con == [4] + + inter1 = kin_explicit_rhs + inter2 = dyn_implicit_rhs + assert simplify(symsystem3.kin_explicit_rhs - inter1) == zeros(2, 1) + assert simplify(symsystem3.dyn_implicit_mat - dyn_implicit_mat) == zeros(3) + assert simplify(symsystem3.dyn_implicit_rhs - inter2) == zeros(3, 1) + + inter = comb_implicit_rhs + assert simplify(symsystem3.comb_implicit_rhs - inter) == zeros(5, 1) + assert simplify(symsystem3.comb_implicit_mat-comb_implicit_mat) == zeros(5) + + inter = comb_explicit_rhs + symsystem3.compute_explicit_form() + assert simplify(symsystem3.comb_explicit_rhs - inter) == zeros(5, 1) + + assert set(symsystem3.dynamic_symbols()) == {y, v, lam, u, x} + assert type(symsystem3.dynamic_symbols()) == tuple + assert set(symsystem3.constant_symbols()) == {l, g, m} + assert type(symsystem3.constant_symbols()) == tuple + + assert symsystem3.output_eqns == {} + + assert symsystem3.bodies == (Pa,) + assert symsystem3.loads == ((P, g * m * N.x),) + + +def test_property_attributes(): + symsystem = SymbolicSystem(states, comb_explicit_rhs, + alg_con=alg_con_full, output_eqns=out_eqns, + coord_idxs=coord_idxs, speed_idxs=speed_idxs, + bodies=bodies, loads=loads) + + with raises(AttributeError): + symsystem.bodies = 42 + with raises(AttributeError): + symsystem.coordinates = 42 + with raises(AttributeError): + symsystem.dyn_implicit_rhs = 42 + with raises(AttributeError): + symsystem.comb_implicit_rhs = 42 + with raises(AttributeError): + symsystem.loads = 42 + with raises(AttributeError): + symsystem.dyn_implicit_mat = 42 + with raises(AttributeError): + symsystem.comb_implicit_mat = 42 + with raises(AttributeError): + symsystem.kin_explicit_rhs = 42 + with raises(AttributeError): + symsystem.comb_explicit_rhs = 42 + with raises(AttributeError): + symsystem.speeds = 42 + with raises(AttributeError): + symsystem.states = 42 + with raises(AttributeError): + symsystem.alg_con = 42 + + +def test_not_specified_errors(): + """This test will cover errors that arise from trying to access attributes + that were not specified upon object creation or were specified on creation + and the user tries to recalculate them.""" + # Trying to access form 2 when form 1 given + # Trying to access form 3 when form 2 given + + symsystem1 = SymbolicSystem(states, comb_explicit_rhs) + + with raises(AttributeError): + symsystem1.comb_implicit_mat + with raises(AttributeError): + symsystem1.comb_implicit_rhs + with raises(AttributeError): + symsystem1.dyn_implicit_mat + with raises(AttributeError): + symsystem1.dyn_implicit_rhs + with raises(AttributeError): + symsystem1.kin_explicit_rhs + with raises(AttributeError): + symsystem1.compute_explicit_form() + + symsystem2 = SymbolicSystem(coordinates, comb_implicit_rhs, speeds=speeds, + mass_matrix=comb_implicit_mat) + + with raises(AttributeError): + symsystem2.dyn_implicit_mat + with raises(AttributeError): + symsystem2.dyn_implicit_rhs + with raises(AttributeError): + symsystem2.kin_explicit_rhs + + # Attribute error when trying to access coordinates and speeds when only the + # states were given. + with raises(AttributeError): + symsystem1.coordinates + with raises(AttributeError): + symsystem1.speeds + + # Attribute error when trying to access bodies and loads when they are not + # given + with raises(AttributeError): + symsystem1.bodies + with raises(AttributeError): + symsystem1.loads + + # Attribute error when trying to access comb_explicit_rhs before it was + # calculated + with raises(AttributeError): + symsystem2.comb_explicit_rhs diff --git a/lib/python3.12/site-packages/sympy/physics/mechanics/tests/test_system_class.py b/lib/python3.12/site-packages/sympy/physics/mechanics/tests/test_system_class.py new file mode 100644 index 0000000000000000000000000000000000000000..924cb8272c27c4f978aa4c3b1999f6ac56e47335 --- /dev/null +++ b/lib/python3.12/site-packages/sympy/physics/mechanics/tests/test_system_class.py @@ -0,0 +1,831 @@ +import pytest + +from sympy.core.symbol import symbols +from sympy.core.sympify import sympify +from sympy.functions.elementary.trigonometric import cos, sin +from sympy.matrices.dense import eye, zeros +from sympy.matrices.immutable import ImmutableMatrix +from sympy.physics.mechanics import ( + Force, KanesMethod, LagrangesMethod, Particle, PinJoint, Point, + PrismaticJoint, ReferenceFrame, RigidBody, Torque, TorqueActuator, System, + dynamicsymbols) +from sympy.simplify.simplify import simplify +from sympy.solvers.solvers import solve + +t = dynamicsymbols._t # type: ignore +q = dynamicsymbols('q:6') # type: ignore +qd = dynamicsymbols('q:6', 1) # type: ignore +u = dynamicsymbols('u:6') # type: ignore +ua = dynamicsymbols('ua:3') # type: ignore + + +class TestSystemBase: + @pytest.fixture() + def _empty_system_setup(self): + self.system = System(ReferenceFrame('frame'), Point('fixed_point')) + + def _empty_system_check(self, exclude=()): + matrices = ('q_ind', 'q_dep', 'q', 'u_ind', 'u_dep', 'u', 'u_aux', + 'kdes', 'holonomic_constraints', 'nonholonomic_constraints') + tuples = ('loads', 'bodies', 'joints', 'actuators') + for attr in matrices: + if attr not in exclude: + assert getattr(self.system, attr)[:] == [] + for attr in tuples: + if attr not in exclude: + assert getattr(self.system, attr) == () + if 'eom_method' not in exclude: + assert self.system.eom_method is None + + def _create_filled_system(self, with_speeds=True): + self.system = System(ReferenceFrame('frame'), Point('fixed_point')) + u = dynamicsymbols('u:6') if with_speeds else qd + self.bodies = symbols('rb1:5', cls=RigidBody) + self.joints = ( + PinJoint('J1', self.bodies[0], self.bodies[1], q[0], u[0]), + PrismaticJoint('J2', self.bodies[1], self.bodies[2], q[1], u[1]), + PinJoint('J3', self.bodies[2], self.bodies[3], q[2], u[2]) + ) + self.system.add_joints(*self.joints) + self.system.add_coordinates(q[3], independent=[False]) + self.system.add_speeds(u[3], independent=False) + if with_speeds: + self.system.add_kdes(u[3] - qd[3]) + self.system.add_auxiliary_speeds(ua[0], ua[1]) + self.system.add_holonomic_constraints(q[2] - q[0] + q[1]) + self.system.add_nonholonomic_constraints(u[3] - qd[1] + u[2]) + self.system.u_ind = u[:2] + self.system.u_dep = u[2:4] + self.q_ind, self.q_dep = self.system.q_ind[:], self.system.q_dep[:] + self.u_ind, self.u_dep = self.system.u_ind[:], self.system.u_dep[:] + self.kdes = self.system.kdes[:] + self.hc = self.system.holonomic_constraints[:] + self.vc = self.system.velocity_constraints[:] + self.nhc = self.system.nonholonomic_constraints[:] + + @pytest.fixture() + def _filled_system_setup(self): + self._create_filled_system(with_speeds=True) + + @pytest.fixture() + def _filled_system_setup_no_speeds(self): + self._create_filled_system(with_speeds=False) + + def _filled_system_check(self, exclude=()): + assert 'q_ind' in exclude or self.system.q_ind[:] == q[:3] + assert 'q_dep' in exclude or self.system.q_dep[:] == [q[3]] + assert 'q' in exclude or self.system.q[:] == q[:4] + assert 'u_ind' in exclude or self.system.u_ind[:] == u[:2] + assert 'u_dep' in exclude or self.system.u_dep[:] == u[2:4] + assert 'u' in exclude or self.system.u[:] == u[:4] + assert 'u_aux' in exclude or self.system.u_aux[:] == ua[:2] + assert 'kdes' in exclude or self.system.kdes[:] == [ + ui - qdi for ui, qdi in zip(u[:4], qd[:4])] + assert ('holonomic_constraints' in exclude or + self.system.holonomic_constraints[:] == [q[2] - q[0] + q[1]]) + assert ('nonholonomic_constraints' in exclude or + self.system.nonholonomic_constraints[:] == [u[3] - qd[1] + u[2]] + ) + assert ('velocity_constraints' in exclude or + self.system.velocity_constraints[:] == [ + qd[2] - qd[0] + qd[1], u[3] - qd[1] + u[2]]) + assert ('bodies' in exclude or + self.system.bodies == tuple(self.bodies)) + assert ('joints' in exclude or + self.system.joints == tuple(self.joints)) + + @pytest.fixture() + def _moving_point_mass(self, _empty_system_setup): + self.system.q_ind = q[0] + self.system.u_ind = u[0] + self.system.kdes = u[0] - q[0].diff(t) + p = Particle('p', mass=symbols('m')) + self.system.add_bodies(p) + p.masscenter.set_pos(self.system.fixed_point, q[0] * self.system.x) + + +class TestSystem(TestSystemBase): + def test_empty_system(self, _empty_system_setup): + self._empty_system_check() + self.system.validate_system() + + def test_filled_system(self, _filled_system_setup): + self._filled_system_check() + self.system.validate_system() + + @pytest.mark.parametrize('frame', [None, ReferenceFrame('frame')]) + @pytest.mark.parametrize('fixed_point', [None, Point('fixed_point')]) + def test_init(self, frame, fixed_point): + if fixed_point is None and frame is None: + self.system = System() + else: + self.system = System(frame, fixed_point) + if fixed_point is None: + assert self.system.fixed_point.name == 'inertial_point' + else: + assert self.system.fixed_point == fixed_point + if frame is None: + assert self.system.frame.name == 'inertial_frame' + else: + assert self.system.frame == frame + self._empty_system_check() + assert isinstance(self.system.q_ind, ImmutableMatrix) + assert isinstance(self.system.q_dep, ImmutableMatrix) + assert isinstance(self.system.q, ImmutableMatrix) + assert isinstance(self.system.u_ind, ImmutableMatrix) + assert isinstance(self.system.u_dep, ImmutableMatrix) + assert isinstance(self.system.u, ImmutableMatrix) + assert isinstance(self.system.kdes, ImmutableMatrix) + assert isinstance(self.system.holonomic_constraints, ImmutableMatrix) + assert isinstance(self.system.nonholonomic_constraints, ImmutableMatrix) + + def test_from_newtonian_rigid_body(self): + rb = RigidBody('body') + self.system = System.from_newtonian(rb) + assert self.system.fixed_point == rb.masscenter + assert self.system.frame == rb.frame + self._empty_system_check(exclude=('bodies',)) + self.system.bodies = (rb,) + + def test_from_newtonian_particle(self): + pt = Particle('particle') + with pytest.raises(TypeError): + System.from_newtonian(pt) + + @pytest.mark.parametrize('args, kwargs, exp_q_ind, exp_q_dep, exp_q', [ + (q[:3], {}, q[:3], [], q[:3]), + (q[:3], {'independent': True}, q[:3], [], q[:3]), + (q[:3], {'independent': False}, [], q[:3], q[:3]), + (q[:3], {'independent': [True, False, True]}, [q[0], q[2]], [q[1]], + [q[0], q[2], q[1]]), + ]) + def test_coordinates(self, _empty_system_setup, args, kwargs, + exp_q_ind, exp_q_dep, exp_q): + # Test add_coordinates + self.system.add_coordinates(*args, **kwargs) + assert self.system.q_ind[:] == exp_q_ind + assert self.system.q_dep[:] == exp_q_dep + assert self.system.q[:] == exp_q + self._empty_system_check(exclude=('q_ind', 'q_dep', 'q')) + # Test setter for q_ind and q_dep + self.system.q_ind = exp_q_ind + self.system.q_dep = exp_q_dep + assert self.system.q_ind[:] == exp_q_ind + assert self.system.q_dep[:] == exp_q_dep + assert self.system.q[:] == exp_q + self._empty_system_check(exclude=('q_ind', 'q_dep', 'q')) + + @pytest.mark.parametrize('func', ['add_coordinates', 'add_speeds']) + @pytest.mark.parametrize('args, kwargs', [ + ((q[0], q[5]), {}), + ((u[0], u[5]), {}), + ((q[0],), {'independent': False}), + ((u[0],), {'independent': False}), + ((u[0], q[5]), {}), + ((symbols('a'), q[5]), {}), + ]) + def test_coordinates_speeds_invalid(self, _filled_system_setup, func, args, + kwargs): + with pytest.raises(ValueError): + getattr(self.system, func)(*args, **kwargs) + self._filled_system_check() + + @pytest.mark.parametrize('args, kwargs, exp_u_ind, exp_u_dep, exp_u', [ + (u[:3], {}, u[:3], [], u[:3]), + (u[:3], {'independent': True}, u[:3], [], u[:3]), + (u[:3], {'independent': False}, [], u[:3], u[:3]), + (u[:3], {'independent': [True, False, True]}, [u[0], u[2]], [u[1]], + [u[0], u[2], u[1]]), + ]) + def test_speeds(self, _empty_system_setup, args, kwargs, exp_u_ind, + exp_u_dep, exp_u): + # Test add_speeds + self.system.add_speeds(*args, **kwargs) + assert self.system.u_ind[:] == exp_u_ind + assert self.system.u_dep[:] == exp_u_dep + assert self.system.u[:] == exp_u + self._empty_system_check(exclude=('u_ind', 'u_dep', 'u')) + # Test setter for u_ind and u_dep + self.system.u_ind = exp_u_ind + self.system.u_dep = exp_u_dep + assert self.system.u_ind[:] == exp_u_ind + assert self.system.u_dep[:] == exp_u_dep + assert self.system.u[:] == exp_u + self._empty_system_check(exclude=('u_ind', 'u_dep', 'u')) + + @pytest.mark.parametrize('args, kwargs, exp_u_aux', [ + (ua[:3], {}, ua[:3]), + ]) + def test_auxiliary_speeds(self, _empty_system_setup, args, kwargs, + exp_u_aux): + # Test add_speeds + self.system.add_auxiliary_speeds(*args, **kwargs) + assert self.system.u_aux[:] == exp_u_aux + self._empty_system_check(exclude=('u_aux',)) + # Test setter for u_ind and u_dep + self.system.u_aux = exp_u_aux + assert self.system.u_aux[:] == exp_u_aux + self._empty_system_check(exclude=('u_aux',)) + + @pytest.mark.parametrize('args, kwargs', [ + ((ua[2], q[0]), {}), + ((ua[2], u[1]), {}), + ((ua[0], ua[2]), {}), + ((symbols('a'), ua[2]), {}), + ]) + def test_auxiliary_invalid(self, _filled_system_setup, args, kwargs): + with pytest.raises(ValueError): + self.system.add_auxiliary_speeds(*args, **kwargs) + self._filled_system_check() + + @pytest.mark.parametrize('prop, add_func, args, kwargs', [ + ('q_ind', 'add_coordinates', (q[0],), {}), + ('q_dep', 'add_coordinates', (q[3],), {'independent': False}), + ('u_ind', 'add_speeds', (u[0],), {}), + ('u_dep', 'add_speeds', (u[3],), {'independent': False}), + ('u_aux', 'add_auxiliary_speeds', (ua[2],), {}), + ('kdes', 'add_kdes', (qd[0] - u[0],), {}), + ('holonomic_constraints', 'add_holonomic_constraints', + (q[0] - q[1],), {}), + ('nonholonomic_constraints', 'add_nonholonomic_constraints', + (u[0] - u[1],), {}), + ('bodies', 'add_bodies', (RigidBody('body'),), {}), + ('loads', 'add_loads', (Force(Point('P'), ReferenceFrame('N').x),), {}), + ('actuators', 'add_actuators', (TorqueActuator( + symbols('T'), ReferenceFrame('N').x, ReferenceFrame('A')),), {}), + ]) + def test_add_after_reset(self, _filled_system_setup, prop, add_func, args, + kwargs): + setattr(self.system, prop, ()) + exclude = (prop, 'q', 'u') + if prop in ('holonomic_constraints', 'nonholonomic_constraints'): + exclude += ('velocity_constraints',) + self._filled_system_check(exclude=exclude) + assert list(getattr(self.system, prop)[:]) == [] + getattr(self.system, add_func)(*args, **kwargs) + assert list(getattr(self.system, prop)[:]) == list(args) + + @pytest.mark.parametrize('prop, add_func, value, error', [ + ('q_ind', 'add_coordinates', symbols('a'), ValueError), + ('q_dep', 'add_coordinates', symbols('a'), ValueError), + ('u_ind', 'add_speeds', symbols('a'), ValueError), + ('u_dep', 'add_speeds', symbols('a'), ValueError), + ('u_aux', 'add_auxiliary_speeds', symbols('a'), ValueError), + ('kdes', 'add_kdes', 7, TypeError), + ('holonomic_constraints', 'add_holonomic_constraints', 7, TypeError), + ('nonholonomic_constraints', 'add_nonholonomic_constraints', 7, + TypeError), + ('bodies', 'add_bodies', symbols('a'), TypeError), + ('loads', 'add_loads', symbols('a'), TypeError), + ('actuators', 'add_actuators', symbols('a'), TypeError), + ]) + def test_type_error(self, _filled_system_setup, prop, add_func, value, + error): + with pytest.raises(error): + getattr(self.system, add_func)(value) + with pytest.raises(error): + setattr(self.system, prop, value) + self._filled_system_check() + + @pytest.mark.parametrize('args, kwargs, exp_kdes', [ + ((), {}, [ui - qdi for ui, qdi in zip(u[:4], qd[:4])]), + ((u[4] - qd[4], u[5] - qd[5]), {}, + [ui - qdi for ui, qdi in zip(u[:6], qd[:6])]), + ]) + def test_kdes(self, _filled_system_setup, args, kwargs, exp_kdes): + # Test add_speeds + self.system.add_kdes(*args, **kwargs) + self._filled_system_check(exclude=('kdes',)) + assert self.system.kdes[:] == exp_kdes + # Test setter for kdes + self.system.kdes = exp_kdes + self._filled_system_check(exclude=('kdes',)) + assert self.system.kdes[:] == exp_kdes + + @pytest.mark.parametrize('args, kwargs', [ + ((u[0] - qd[0], u[4] - qd[4]), {}), + ((-(u[0] - qd[0]), u[4] - qd[4]), {}), + (([u[0] - u[0], u[4] - qd[4]]), {}), + ]) + def test_kdes_invalid(self, _filled_system_setup, args, kwargs): + with pytest.raises(ValueError): + self.system.add_kdes(*args, **kwargs) + self._filled_system_check() + + @pytest.mark.parametrize('args, kwargs, exp_con', [ + ((), {}, [q[2] - q[0] + q[1]]), + ((q[4] - q[5], q[5] + q[3]), {}, + [q[2] - q[0] + q[1], q[4] - q[5], q[5] + q[3]]), + ]) + def test_holonomic_constraints(self, _filled_system_setup, args, kwargs, + exp_con): + exclude = ('holonomic_constraints', 'velocity_constraints') + exp_vel_con = [c.diff(t) for c in exp_con] + self.nhc + # Test add_holonomic_constraints + self.system.add_holonomic_constraints(*args, **kwargs) + self._filled_system_check(exclude=exclude) + assert self.system.holonomic_constraints[:] == exp_con + assert self.system.velocity_constraints[:] == exp_vel_con + # Test setter for holonomic_constraints + self.system.holonomic_constraints = exp_con + self._filled_system_check(exclude=exclude) + assert self.system.holonomic_constraints[:] == exp_con + assert self.system.velocity_constraints[:] == exp_vel_con + + @pytest.mark.parametrize('args, kwargs', [ + ((q[2] - q[0] + q[1], q[4] - q[3]), {}), + ((-(q[2] - q[0] + q[1]), q[4] - q[3]), {}), + ((q[0] - q[0], q[4] - q[3]), {}), + ]) + def test_holonomic_constraints_invalid(self, _filled_system_setup, args, + kwargs): + with pytest.raises(ValueError): + self.system.add_holonomic_constraints(*args, **kwargs) + self._filled_system_check() + + @pytest.mark.parametrize('args, kwargs, exp_con', [ + ((), {}, [u[3] - qd[1] + u[2]]), + ((u[4] - u[5], u[5] + u[3]), {}, + [u[3] - qd[1] + u[2], u[4] - u[5], u[5] + u[3]]), + ]) + def test_nonholonomic_constraints(self, _filled_system_setup, args, kwargs, + exp_con): + exclude = ('nonholonomic_constraints', 'velocity_constraints') + exp_vel_con = self.vc[:len(self.hc)] + exp_con + # Test add_nonholonomic_constraints + self.system.add_nonholonomic_constraints(*args, **kwargs) + self._filled_system_check(exclude=exclude) + assert self.system.nonholonomic_constraints[:] == exp_con + assert self.system.velocity_constraints[:] == exp_vel_con + # Test setter for nonholonomic_constraints + self.system.nonholonomic_constraints = exp_con + self._filled_system_check(exclude=exclude) + assert self.system.nonholonomic_constraints[:] == exp_con + assert self.system.velocity_constraints[:] == exp_vel_con + + @pytest.mark.parametrize('args, kwargs', [ + ((u[3] - qd[1] + u[2], u[4] - u[3]), {}), + ((-(u[3] - qd[1] + u[2]), u[4] - u[3]), {}), + ((u[0] - u[0], u[4] - u[3]), {}), + (([u[0] - u[0], u[4] - u[3]]), {}), + ]) + def test_nonholonomic_constraints_invalid(self, _filled_system_setup, args, + kwargs): + with pytest.raises(ValueError): + self.system.add_nonholonomic_constraints(*args, **kwargs) + self._filled_system_check() + + @pytest.mark.parametrize('constraints, expected', [ + ([], []), + (qd[2] - qd[0] + qd[1], [qd[2] - qd[0] + qd[1]]), + ([qd[2] + qd[1], u[2] - u[1]], [qd[2] + qd[1], u[2] - u[1]]), + ]) + def test_velocity_constraints_overwrite(self, _filled_system_setup, + constraints, expected): + self.system.velocity_constraints = constraints + self._filled_system_check(exclude=('velocity_constraints',)) + assert self.system.velocity_constraints[:] == expected + + def test_velocity_constraints_back_to_auto(self, _filled_system_setup): + self.system.velocity_constraints = qd[3] - qd[2] + self._filled_system_check(exclude=('velocity_constraints',)) + assert self.system.velocity_constraints[:] == [qd[3] - qd[2]] + self.system.velocity_constraints = None + self._filled_system_check() + + def test_bodies(self, _filled_system_setup): + rb1, rb2 = RigidBody('rb1'), RigidBody('rb2') + p1, p2 = Particle('p1'), Particle('p2') + self.system.add_bodies(rb1, p1) + assert self.system.bodies == (*self.bodies, rb1, p1) + self.system.add_bodies(p2) + assert self.system.bodies == (*self.bodies, rb1, p1, p2) + self.system.bodies = [] + assert self.system.bodies == () + self.system.bodies = p2 + assert self.system.bodies == (p2,) + symb = symbols('symb') + pytest.raises(TypeError, lambda: self.system.add_bodies(symb)) + pytest.raises(ValueError, lambda: self.system.add_bodies(p2)) + with pytest.raises(TypeError): + self.system.bodies = (rb1, rb2, p1, p2, symb) + assert self.system.bodies == (p2,) + + def test_add_loads(self): + system = System() + N, A = ReferenceFrame('N'), ReferenceFrame('A') + rb1 = RigidBody('rb1', frame=N) + mc1 = Point('mc1') + p1 = Particle('p1', mc1) + system.add_loads(Torque(rb1, N.x), (mc1, A.x), Force(p1, A.x)) + assert system.loads == ((N, N.x), (mc1, A.x), (mc1, A.x)) + system.loads = [(A, A.x)] + assert system.loads == ((A, A.x),) + pytest.raises(ValueError, lambda: system.add_loads((N, N.x, N.y))) + with pytest.raises(TypeError): + system.loads = (N, N.x) + assert system.loads == ((A, A.x),) + + def test_add_actuators(self): + system = System() + N, A = ReferenceFrame('N'), ReferenceFrame('A') + act1 = TorqueActuator(symbols('T1'), N.x, N) + act2 = TorqueActuator(symbols('T2'), N.y, N, A) + system.add_actuators(act1) + assert system.actuators == (act1,) + assert system.loads == () + system.actuators = (act2,) + assert system.actuators == (act2,) + + def test_add_joints(self): + q1, q2, q3, q4, u1, u2, u3 = dynamicsymbols('q1:5 u1:4') + rb1, rb2, rb3, rb4, rb5 = symbols('rb1:6', cls=RigidBody) + J1 = PinJoint('J1', rb1, rb2, q1, u1) + J2 = PrismaticJoint('J2', rb2, rb3, q2, u2) + J3 = PinJoint('J3', rb3, rb4, q3, u3) + J_lag = PinJoint('J_lag', rb4, rb5, q4, q4.diff(t)) + system = System() + system.add_joints(J1) + assert system.joints == (J1,) + assert system.bodies == (rb1, rb2) + assert system.q_ind == ImmutableMatrix([q1]) + assert system.u_ind == ImmutableMatrix([u1]) + assert system.kdes == ImmutableMatrix([u1 - q1.diff(t)]) + system.add_bodies(rb4) + system.add_coordinates(q3) + system.add_kdes(u3 - q3.diff(t)) + system.add_joints(J3) + assert system.joints == (J1, J3) + assert system.bodies == (rb1, rb2, rb4, rb3) + assert system.q_ind == ImmutableMatrix([q1, q3]) + assert system.u_ind == ImmutableMatrix([u1, u3]) + assert system.kdes == ImmutableMatrix( + [u1 - q1.diff(t), u3 - q3.diff(t)]) + system.add_kdes(-(u2 - q2.diff(t))) + system.add_joints(J2) + assert system.joints == (J1, J3, J2) + assert system.bodies == (rb1, rb2, rb4, rb3) + assert system.q_ind == ImmutableMatrix([q1, q3, q2]) + assert system.u_ind == ImmutableMatrix([u1, u3, u2]) + assert system.kdes == ImmutableMatrix([u1 - q1.diff(t), u3 - q3.diff(t), + -(u2 - q2.diff(t))]) + system.add_joints(J_lag) + assert system.joints == (J1, J3, J2, J_lag) + assert system.bodies == (rb1, rb2, rb4, rb3, rb5) + assert system.q_ind == ImmutableMatrix([q1, q3, q2, q4]) + assert system.u_ind == ImmutableMatrix([u1, u3, u2, q4.diff(t)]) + assert system.kdes == ImmutableMatrix([u1 - q1.diff(t), u3 - q3.diff(t), + -(u2 - q2.diff(t))]) + assert system.q_dep[:] == [] + assert system.u_dep[:] == [] + pytest.raises(ValueError, lambda: system.add_joints(J2)) + pytest.raises(TypeError, lambda: system.add_joints(rb1)) + + def test_joints_setter(self, _filled_system_setup): + self.system.joints = self.joints[1:] + assert self.system.joints == self.joints[1:] + self._filled_system_check(exclude=('joints',)) + self.system.q_ind = () + self.system.u_ind = () + self.system.joints = self.joints + self._filled_system_check() + + @pytest.mark.parametrize('name, joint_index', [ + ('J1', 0), + ('J2', 1), + ('not_existing', None), + ]) + def test_get_joint(self, _filled_system_setup, name, joint_index): + joint = self.system.get_joint(name) + if joint_index is None: + assert joint is None + else: + assert joint == self.joints[joint_index] + + @pytest.mark.parametrize('name, body_index', [ + ('rb1', 0), + ('rb3', 2), + ('not_existing', None), + ]) + def test_get_body(self, _filled_system_setup, name, body_index): + body = self.system.get_body(name) + if body_index is None: + assert body is None + else: + assert body == self.bodies[body_index] + + @pytest.mark.parametrize('eom_method', [KanesMethod, LagrangesMethod]) + def test_form_eoms_calls_subclass(self, _moving_point_mass, eom_method): + class MyMethod(eom_method): + pass + + self.system.form_eoms(eom_method=MyMethod) + assert isinstance(self.system.eom_method, MyMethod) + + @pytest.mark.parametrize('kwargs, expected', [ + ({}, ImmutableMatrix([[-1, 0], [0, symbols('m')]])), + ({'explicit_kinematics': True}, ImmutableMatrix([[1, 0], + [0, symbols('m')]])), + ]) + def test_system_kane_form_eoms_kwargs(self, _moving_point_mass, kwargs, + expected): + self.system.form_eoms(**kwargs) + assert self.system.mass_matrix_full == expected + + @pytest.mark.parametrize('kwargs, mm, gm', [ + ({}, ImmutableMatrix([[1, 0], [0, symbols('m')]]), + ImmutableMatrix([q[0].diff(t), 0])), + ]) + def test_system_lagrange_form_eoms_kwargs(self, _moving_point_mass, kwargs, + mm, gm): + self.system.form_eoms(eom_method=LagrangesMethod, **kwargs) + assert self.system.mass_matrix_full == mm + assert self.system.forcing_full == gm + + @pytest.mark.parametrize('eom_method, kwargs, error', [ + (KanesMethod, {'non_existing_kwarg': 1}, TypeError), + (LagrangesMethod, {'non_existing_kwarg': 1}, TypeError), + (KanesMethod, {'bodies': []}, ValueError), + (KanesMethod, {'kd_eqs': []}, ValueError), + (LagrangesMethod, {'bodies': []}, ValueError), + (LagrangesMethod, {'Lagrangian': 1}, ValueError), + ]) + def test_form_eoms_kwargs_errors(self, _empty_system_setup, eom_method, + kwargs, error): + self.system.q_ind = q[0] + p = Particle('p', mass=symbols('m')) + self.system.add_bodies(p) + p.masscenter.set_pos(self.system.fixed_point, q[0] * self.system.x) + with pytest.raises(error): + self.system.form_eoms(eom_method=eom_method, **kwargs) + + +class TestValidateSystem(TestSystemBase): + @pytest.mark.parametrize('valid_method, invalid_method, with_speeds', [ + (KanesMethod, LagrangesMethod, True), + (LagrangesMethod, KanesMethod, False) + ]) + def test_only_valid(self, valid_method, invalid_method, with_speeds): + self._create_filled_system(with_speeds=with_speeds) + self.system.validate_system(valid_method) + # Test Lagrange should fail due to the usage of generalized speeds + with pytest.raises(ValueError): + self.system.validate_system(invalid_method) + + @pytest.mark.parametrize('method, with_speeds', [ + (KanesMethod, True), (LagrangesMethod, False)]) + def test_missing_joint_coordinate(self, method, with_speeds): + self._create_filled_system(with_speeds=with_speeds) + self.system.q_ind = self.q_ind[1:] + self.system.u_ind = self.u_ind[:-1] + self.system.kdes = self.kdes[:-1] + pytest.raises(ValueError, lambda: self.system.validate_system(method)) + + def test_missing_joint_speed(self, _filled_system_setup): + self.system.q_ind = self.q_ind[:-1] + self.system.u_ind = self.u_ind[1:] + self.system.kdes = self.kdes[:-1] + pytest.raises(ValueError, lambda: self.system.validate_system()) + + def test_missing_joint_kdes(self, _filled_system_setup): + self.system.kdes = self.kdes[1:] + pytest.raises(ValueError, lambda: self.system.validate_system()) + + def test_negative_joint_kdes(self, _filled_system_setup): + self.system.kdes = [-self.kdes[0]] + self.kdes[1:] + self.system.validate_system() + + @pytest.mark.parametrize('method, with_speeds', [ + (KanesMethod, True), (LagrangesMethod, False)]) + def test_missing_holonomic_constraint(self, method, with_speeds): + self._create_filled_system(with_speeds=with_speeds) + self.system.holonomic_constraints = [] + self.system.nonholonomic_constraints = self.nhc + [ + self.u_ind[1] - self.u_dep[0] + self.u_ind[0]] + pytest.raises(ValueError, lambda: self.system.validate_system(method)) + self.system.q_dep = [] + self.system.q_ind = self.q_ind + self.q_dep + self.system.validate_system(method) + + def test_missing_nonholonomic_constraint(self, _filled_system_setup): + self.system.nonholonomic_constraints = [] + pytest.raises(ValueError, lambda: self.system.validate_system()) + self.system.u_dep = self.u_dep[1] + self.system.u_ind = self.u_ind + [self.u_dep[0]] + self.system.validate_system() + + def test_number_of_coordinates_speeds(self, _filled_system_setup): + # Test more speeds than coordinates + self.system.u_ind = self.u_ind + [u[5]] + self.system.kdes = self.kdes + [u[5] - qd[5]] + self.system.validate_system() + # Test more coordinates than speeds + self.system.q_ind = self.q_ind + self.system.u_ind = self.u_ind[:-1] + self.system.kdes = self.kdes[:-1] + pytest.raises(ValueError, lambda: self.system.validate_system()) + + def test_number_of_kdes(self, _filled_system_setup): + # Test wrong number of kdes + self.system.kdes = self.kdes[:-1] + pytest.raises(ValueError, lambda: self.system.validate_system()) + self.system.kdes = self.kdes + [u[2] + u[1] - qd[2]] + pytest.raises(ValueError, lambda: self.system.validate_system()) + + def test_duplicates(self, _filled_system_setup): + # This is basically a redundant feature, which should never fail + self.system.validate_system(check_duplicates=True) + + def test_speeds_in_lagrange(self, _filled_system_setup_no_speeds): + self.system.u_ind = u[:len(self.u_ind)] + with pytest.raises(ValueError): + self.system.validate_system(LagrangesMethod) + self.system.u_ind = [] + self.system.validate_system(LagrangesMethod) + self.system.u_aux = ua + with pytest.raises(ValueError): + self.system.validate_system(LagrangesMethod) + self.system.u_aux = [] + self.system.validate_system(LagrangesMethod) + self.system.add_joints( + PinJoint('Ju', RigidBody('rbu1'), RigidBody('rbu2'))) + self.system.u_ind = [] + with pytest.raises(ValueError): + self.system.validate_system(LagrangesMethod) + + +class TestSystemExamples: + def test_cart_pendulum_kanes(self): + # This example is the same as in the top documentation of System + # Added a spring to the cart + g, l, mc, mp, k = symbols('g l mc mp k') + F, qp, qc, up, uc = dynamicsymbols('F qp qc up uc') + rail = RigidBody('rail') + cart = RigidBody('cart', mass=mc) + bob = Particle('bob', mass=mp) + bob_frame = ReferenceFrame('bob_frame') + system = System.from_newtonian(rail) + assert system.bodies == (rail,) + assert system.frame == rail.frame + assert system.fixed_point == rail.masscenter + slider = PrismaticJoint('slider', rail, cart, qc, uc, joint_axis=rail.x) + pin = PinJoint('pin', cart, bob, qp, up, joint_axis=cart.z, + child_interframe=bob_frame, child_point=l * bob_frame.y) + system.add_joints(slider, pin) + assert system.joints == (slider, pin) + assert system.get_joint('slider') == slider + assert system.get_body('bob') == bob + system.apply_uniform_gravity(-g * system.y) + system.add_loads((cart.masscenter, F * rail.x)) + system.add_actuators(TorqueActuator(k * qp, cart.z, bob_frame, cart)) + system.validate_system() + system.form_eoms() + assert isinstance(system.eom_method, KanesMethod) + assert (simplify(system.mass_matrix - ImmutableMatrix( + [[mp + mc, mp * l * cos(qp)], [mp * l * cos(qp), mp * l ** 2]])) + == zeros(2, 2)) + assert (simplify(system.forcing - ImmutableMatrix([ + [mp * l * up ** 2 * sin(qp) + F], + [-mp * g * l * sin(qp) + k * qp]])) == zeros(2, 1)) + + system.add_holonomic_constraints( + sympify(bob.masscenter.pos_from(rail.masscenter).dot(system.x))) + assert system.eom_method is None + system.q_ind, system.q_dep = qp, qc + system.u_ind, system.u_dep = up, uc + system.validate_system() + + # Computed solution based on manually solving the constraints + subs = {qc: -l * sin(qp), + uc: -l * cos(qp) * up, + uc.diff(t): l * (up ** 2 * sin(qp) - up.diff(t) * cos(qp))} + upd_expected = ( + (-g * mp * sin(qp) + k * qp / l + l * mc * sin(2 * qp) * up ** 2 / 2 + - l * mp * sin(2 * qp) * up ** 2 / 2 - F * cos(qp)) / + (l * (mc * cos(qp) ** 2 + mp * sin(qp) ** 2))) + upd_sol = tuple(solve(system.form_eoms().xreplace(subs), + up.diff(t)).values())[0] + assert simplify(upd_sol - upd_expected) == 0 + assert isinstance(system.eom_method, KanesMethod) + + # Test other output + Mk = -ImmutableMatrix([[0, 1], [1, 0]]) + gk = -ImmutableMatrix([uc, up]) + Md = ImmutableMatrix([[-l ** 2 * mp * cos(qp) ** 2 + l ** 2 * mp, + l * mp * cos(qp) - l * (mc + mp) * cos(qp)], + [l * cos(qp), 1]]) + gd = ImmutableMatrix( + [[-g * l * mp * sin(qp) + k * qp - l ** 2 * mp * up ** 2 * sin(qp) * + cos(qp) - l * F * cos(qp)], [l * up ** 2 * sin(qp)]]) + Mm = (Mk.row_join(zeros(2, 2))).col_join(zeros(2, 2).row_join(Md)) + gm = gk.col_join(gd) + assert simplify(system.mass_matrix - Md) == zeros(2, 2) + assert simplify(system.forcing - gd) == zeros(2, 1) + assert simplify(system.mass_matrix_full - Mm) == zeros(4, 4) + assert simplify(system.forcing_full - gm) == zeros(4, 1) + + def test_cart_pendulum_lagrange(self): + # Lagrange version of test_cart_pendulus_kanes + # Added a spring to the cart + g, l, mc, mp, k = symbols('g l mc mp k') + F, qp, qc = dynamicsymbols('F qp qc') + qpd, qcd = dynamicsymbols('qp qc', 1) + rail = RigidBody('rail') + cart = RigidBody('cart', mass=mc) + bob = Particle('bob', mass=mp) + bob_frame = ReferenceFrame('bob_frame') + system = System.from_newtonian(rail) + assert system.bodies == (rail,) + assert system.frame == rail.frame + assert system.fixed_point == rail.masscenter + slider = PrismaticJoint('slider', rail, cart, qc, qcd, + joint_axis=rail.x) + pin = PinJoint('pin', cart, bob, qp, qpd, joint_axis=cart.z, + child_interframe=bob_frame, child_point=l * bob_frame.y) + system.add_joints(slider, pin) + assert system.joints == (slider, pin) + assert system.get_joint('slider') == slider + assert system.get_body('bob') == bob + for body in system.bodies: + body.potential_energy = body.mass * g * body.masscenter.pos_from( + system.fixed_point).dot(system.y) + system.add_loads((cart.masscenter, F * rail.x)) + system.add_actuators(TorqueActuator(k * qp, cart.z, bob_frame, cart)) + system.validate_system(LagrangesMethod) + system.form_eoms(LagrangesMethod) + assert (simplify(system.mass_matrix - ImmutableMatrix( + [[mp + mc, mp * l * cos(qp)], [mp * l * cos(qp), mp * l ** 2]])) + == zeros(2, 2)) + assert (simplify(system.forcing - ImmutableMatrix([ + [mp * l * qpd ** 2 * sin(qp) + F], [-mp * g * l * sin(qp) + k * qp]] + )) == zeros(2, 1)) + + system.add_holonomic_constraints( + sympify(bob.masscenter.pos_from(rail.masscenter).dot(system.x))) + assert system.eom_method is None + system.q_ind, system.q_dep = qp, qc + + # Computed solution based on manually solving the constraints + subs = {qc: -l * sin(qp), + qcd: -l * cos(qp) * qpd, + qcd.diff(t): l * (qpd ** 2 * sin(qp) - qpd.diff(t) * cos(qp))} + qpdd_expected = ( + (-g * mp * sin(qp) + k * qp / l + l * mc * sin(2 * qp) * qpd ** 2 / + 2 - l * mp * sin(2 * qp) * qpd ** 2 / 2 - F * cos(qp)) / + (l * (mc * cos(qp) ** 2 + mp * sin(qp) ** 2))) + eoms = system.form_eoms(LagrangesMethod) + lam1 = system.eom_method.lam_vec[0] + lam1_sol = system.eom_method.solve_multipliers()[lam1] + qpdd_sol = solve(eoms[0].xreplace({lam1: lam1_sol}).xreplace(subs), + qpd.diff(t))[0] + assert simplify(qpdd_sol - qpdd_expected) == 0 + assert isinstance(system.eom_method, LagrangesMethod) + + # Test other output + Md = ImmutableMatrix([[l ** 2 * mp, l * mp * cos(qp), -l * cos(qp)], + [l * mp * cos(qp), mc + mp, -1]]) + gd = ImmutableMatrix( + [[-g * l * mp * sin(qp) + k * qp], + [l * mp * sin(qp) * qpd ** 2 + F]]) + Mm = (eye(2).row_join(zeros(2, 3))).col_join(zeros(3, 2).row_join( + Md.col_join(ImmutableMatrix([l * cos(qp), 1, 0]).T))) + gm = ImmutableMatrix([qpd, qcd] + gd[:] + [l * sin(qp) * qpd ** 2]) + assert simplify(system.mass_matrix - Md) == zeros(2, 3) + assert simplify(system.forcing - gd) == zeros(2, 1) + assert simplify(system.mass_matrix_full - Mm) == zeros(5, 5) + assert simplify(system.forcing_full - gm) == zeros(5, 1) + + def test_box_on_ground(self): + # Particle sliding on ground with friction. The applied force is assumed + # to be positive and to be higher than the friction force. + g, m, mu = symbols('g m mu') + q, u, ua = dynamicsymbols('q u ua') + N, F = dynamicsymbols('N F', positive=True) + P = Particle("P", mass=m) + system = System() + system.add_bodies(P) + P.masscenter.set_pos(system.fixed_point, q * system.x) + P.masscenter.set_vel(system.frame, u * system.x + ua * system.y) + system.q_ind, system.u_ind, system.u_aux = [q], [u], [ua] + system.kdes = [q.diff(t) - u] + system.apply_uniform_gravity(-g * system.y) + system.add_loads( + Force(P, N * system.y), + Force(P, F * system.x - mu * N * system.x)) + system.validate_system() + system.form_eoms() + + # Test other output + Mk = ImmutableMatrix([1]) + gk = ImmutableMatrix([u]) + Md = ImmutableMatrix([m]) + gd = ImmutableMatrix([F - mu * N]) + Mm = (Mk.row_join(zeros(1, 1))).col_join(zeros(1, 1).row_join(Md)) + gm = gk.col_join(gd) + aux_eqs = ImmutableMatrix([N - m * g]) + assert simplify(system.mass_matrix - Md) == zeros(1, 1) + assert simplify(system.forcing - gd) == zeros(1, 1) + assert simplify(system.mass_matrix_full - Mm) == zeros(2, 2) + assert simplify(system.forcing_full - gm) == zeros(2, 1) + assert simplify(system.eom_method.auxiliary_eqs - aux_eqs + ) == zeros(1, 1) diff --git a/lib/python3.12/site-packages/sympy/physics/mechanics/tests/test_wrapping_geometry.py b/lib/python3.12/site-packages/sympy/physics/mechanics/tests/test_wrapping_geometry.py new file mode 100644 index 0000000000000000000000000000000000000000..30c3ae71db5da75238ebb3d4cc53e11a29a72e5d --- /dev/null +++ b/lib/python3.12/site-packages/sympy/physics/mechanics/tests/test_wrapping_geometry.py @@ -0,0 +1,363 @@ +"""Tests for the ``sympy.physics.mechanics.wrapping_geometry.py`` module.""" + +import pytest + +from sympy import ( + Integer, + Rational, + S, + Symbol, + acos, + cos, + pi, + sin, + sqrt, +) +from sympy.core.relational import Eq +from sympy.physics.mechanics import ( + Point, + ReferenceFrame, + WrappingCylinder, + WrappingSphere, + dynamicsymbols, +) +from sympy.simplify.simplify import simplify + + +r = Symbol('r', positive=True) +x = Symbol('x') +q = dynamicsymbols('q') +N = ReferenceFrame('N') + + +class TestWrappingSphere: + + @staticmethod + def test_valid_constructor(): + r = Symbol('r', positive=True) + pO = Point('pO') + sphere = WrappingSphere(r, pO) + assert isinstance(sphere, WrappingSphere) + assert hasattr(sphere, 'radius') + assert sphere.radius == r + assert hasattr(sphere, 'point') + assert sphere.point == pO + + @staticmethod + @pytest.mark.parametrize('position', [S.Zero, Integer(2)*r*N.x]) + def test_geodesic_length_point_not_on_surface_invalid(position): + r = Symbol('r', positive=True) + pO = Point('pO') + sphere = WrappingSphere(r, pO) + + p1 = Point('p1') + p1.set_pos(pO, position) + p2 = Point('p2') + p2.set_pos(pO, position) + + error_msg = r'point .* does not lie on the surface of' + with pytest.raises(ValueError, match=error_msg): + sphere.geodesic_length(p1, p2) + + @staticmethod + @pytest.mark.parametrize( + 'position_1, position_2, expected', + [ + (r*N.x, r*N.x, S.Zero), + (r*N.x, r*N.y, S.Half*pi*r), + (r*N.x, r*-N.x, pi*r), + (r*-N.x, r*N.x, pi*r), + (r*N.x, r*sqrt(2)*S.Half*(N.x + N.y), Rational(1, 4)*pi*r), + ( + r*sqrt(2)*S.Half*(N.x + N.y), + r*sqrt(3)*Rational(1, 3)*(N.x + N.y + N.z), + r*acos(sqrt(6)*Rational(1, 3)), + ), + ] + ) + def test_geodesic_length(position_1, position_2, expected): + r = Symbol('r', positive=True) + pO = Point('pO') + sphere = WrappingSphere(r, pO) + + p1 = Point('p1') + p1.set_pos(pO, position_1) + p2 = Point('p2') + p2.set_pos(pO, position_2) + + assert simplify(Eq(sphere.geodesic_length(p1, p2), expected)) + + @staticmethod + @pytest.mark.parametrize( + 'position_1, position_2, vector_1, vector_2', + [ + (r * N.x, r * N.y, N.y, N.x), + (r * N.x, -r * N.y, -N.y, N.x), + ( + r * N.y, + sqrt(2)/2 * r * N.x - sqrt(2)/2 * r * N.y, + N.x, + sqrt(2)/2 * N.x + sqrt(2)/2 * N.y, + ), + ( + r * N.x, + r / 2 * N.x + sqrt(3)/2 * r * N.y, + N.y, + sqrt(3)/2 * N.x - 1/2 * N.y, + ), + ( + r * N.x, + sqrt(2)/2 * r * N.x + sqrt(2)/2 * r * N.y, + N.y, + sqrt(2)/2 * N.x - sqrt(2)/2 * N.y, + ), + ] + ) + def test_geodesic_end_vectors(position_1, position_2, vector_1, vector_2): + r = Symbol('r', positive=True) + pO = Point('pO') + sphere = WrappingSphere(r, pO) + + p1 = Point('p1') + p1.set_pos(pO, position_1) + p2 = Point('p2') + p2.set_pos(pO, position_2) + + expected = (vector_1, vector_2) + + assert sphere.geodesic_end_vectors(p1, p2) == expected + + @staticmethod + @pytest.mark.parametrize( + 'position', + [r * N.x, r * cos(q) * N.x + r * sin(q) * N.y] + ) + def test_geodesic_end_vectors_invalid_coincident(position): + r = Symbol('r', positive=True) + pO = Point('pO') + sphere = WrappingSphere(r, pO) + + p1 = Point('p1') + p1.set_pos(pO, position) + p2 = Point('p2') + p2.set_pos(pO, position) + + with pytest.raises(ValueError): + _ = sphere.geodesic_end_vectors(p1, p2) + + @staticmethod + @pytest.mark.parametrize( + 'position_1, position_2', + [ + (r * N.x, -r * N.x), + (-r * N.y, r * N.y), + ( + r * cos(q) * N.x + r * sin(q) * N.y, + -r * cos(q) * N.x - r * sin(q) * N.y, + ) + ] + ) + def test_geodesic_end_vectors_invalid_diametrically_opposite( + position_1, + position_2, + ): + r = Symbol('r', positive=True) + pO = Point('pO') + sphere = WrappingSphere(r, pO) + + p1 = Point('p1') + p1.set_pos(pO, position_1) + p2 = Point('p2') + p2.set_pos(pO, position_2) + + with pytest.raises(ValueError): + _ = sphere.geodesic_end_vectors(p1, p2) + + +class TestWrappingCylinder: + + @staticmethod + def test_valid_constructor(): + N = ReferenceFrame('N') + r = Symbol('r', positive=True) + pO = Point('pO') + cylinder = WrappingCylinder(r, pO, N.x) + assert isinstance(cylinder, WrappingCylinder) + assert hasattr(cylinder, 'radius') + assert cylinder.radius == r + assert hasattr(cylinder, 'point') + assert cylinder.point == pO + assert hasattr(cylinder, 'axis') + assert cylinder.axis == N.x + + @staticmethod + @pytest.mark.parametrize( + 'position, expected', + [ + (S.Zero, False), + (r*N.y, True), + (r*N.z, True), + (r*(N.y + N.z).normalize(), True), + (Integer(2)*r*N.y, False), + (r*(N.x + N.y), True), + (r*(Integer(2)*N.x + N.y), True), + (Integer(2)*N.x + r*(Integer(2)*N.y + N.z).normalize(), True), + (r*(cos(q)*N.y + sin(q)*N.z), True) + ] + ) + def test_point_is_on_surface(position, expected): + r = Symbol('r', positive=True) + pO = Point('pO') + cylinder = WrappingCylinder(r, pO, N.x) + + p1 = Point('p1') + p1.set_pos(pO, position) + + assert cylinder.point_on_surface(p1) is expected + + @staticmethod + @pytest.mark.parametrize('position', [S.Zero, Integer(2)*r*N.y]) + def test_geodesic_length_point_not_on_surface_invalid(position): + r = Symbol('r', positive=True) + pO = Point('pO') + cylinder = WrappingCylinder(r, pO, N.x) + + p1 = Point('p1') + p1.set_pos(pO, position) + p2 = Point('p2') + p2.set_pos(pO, position) + + error_msg = r'point .* does not lie on the surface of' + with pytest.raises(ValueError, match=error_msg): + cylinder.geodesic_length(p1, p2) + + @staticmethod + @pytest.mark.parametrize( + 'axis, position_1, position_2, expected', + [ + (N.x, r*N.y, r*N.y, S.Zero), + (N.x, r*N.y, N.x + r*N.y, S.One), + (N.x, r*N.y, -x*N.x + r*N.y, sqrt(x**2)), + (-N.x, r*N.y, x*N.x + r*N.y, sqrt(x**2)), + (N.x, r*N.y, r*N.z, S.Half*pi*sqrt(r**2)), + (-N.x, r*N.y, r*N.z, Integer(3)*S.Half*pi*sqrt(r**2)), + (N.x, r*N.z, r*N.y, Integer(3)*S.Half*pi*sqrt(r**2)), + (-N.x, r*N.z, r*N.y, S.Half*pi*sqrt(r**2)), + (N.x, r*N.y, r*(cos(q)*N.y + sin(q)*N.z), sqrt(r**2*q**2)), + ( + -N.x, r*N.y, + r*(cos(q)*N.y + sin(q)*N.z), + sqrt(r**2*(Integer(2)*pi - q)**2), + ), + ] + ) + def test_geodesic_length(axis, position_1, position_2, expected): + r = Symbol('r', positive=True) + pO = Point('pO') + cylinder = WrappingCylinder(r, pO, axis) + + p1 = Point('p1') + p1.set_pos(pO, position_1) + p2 = Point('p2') + p2.set_pos(pO, position_2) + + assert simplify(Eq(cylinder.geodesic_length(p1, p2), expected)) + + @staticmethod + @pytest.mark.parametrize( + 'axis, position_1, position_2, vector_1, vector_2', + [ + (N.z, r * N.x, r * N.y, N.y, N.x), + (N.z, r * N.x, -r * N.x, N.y, N.y), + (N.z, -r * N.x, r * N.x, -N.y, -N.y), + (-N.z, r * N.x, -r * N.x, -N.y, -N.y), + (-N.z, -r * N.x, r * N.x, N.y, N.y), + (N.z, r * N.x, -r * N.y, N.y, -N.x), + ( + N.z, + r * N.y, + sqrt(2)/2 * r * N.x - sqrt(2)/2 * r * N.y, + - N.x, + - sqrt(2)/2 * N.x - sqrt(2)/2 * N.y, + ), + ( + N.z, + r * N.x, + r / 2 * N.x + sqrt(3)/2 * r * N.y, + N.y, + sqrt(3)/2 * N.x - 1/2 * N.y, + ), + ( + N.z, + r * N.x, + sqrt(2)/2 * r * N.x + sqrt(2)/2 * r * N.y, + N.y, + sqrt(2)/2 * N.x - sqrt(2)/2 * N.y, + ), + ( + N.z, + r * N.x, + r * N.x + N.z, + N.z, + -N.z, + ), + ( + N.z, + r * N.x, + r * N.y + pi/2 * r * N.z, + sqrt(2)/2 * N.y + sqrt(2)/2 * N.z, + sqrt(2)/2 * N.x - sqrt(2)/2 * N.z, + ), + ( + N.z, + r * N.x, + r * cos(q) * N.x + r * sin(q) * N.y, + N.y, + sin(q) * N.x - cos(q) * N.y, + ), + ] + ) + def test_geodesic_end_vectors( + axis, + position_1, + position_2, + vector_1, + vector_2, + ): + r = Symbol('r', positive=True) + pO = Point('pO') + cylinder = WrappingCylinder(r, pO, axis) + + p1 = Point('p1') + p1.set_pos(pO, position_1) + p2 = Point('p2') + p2.set_pos(pO, position_2) + + expected = (vector_1, vector_2) + end_vectors = tuple( + end_vector.simplify() + for end_vector in cylinder.geodesic_end_vectors(p1, p2) + ) + + assert end_vectors == expected + + @staticmethod + @pytest.mark.parametrize( + 'axis, position', + [ + (N.z, r * N.x), + (N.z, r * cos(q) * N.x + r * sin(q) * N.y + N.z), + ] + ) + def test_geodesic_end_vectors_invalid_coincident(axis, position): + r = Symbol('r', positive=True) + pO = Point('pO') + cylinder = WrappingCylinder(r, pO, axis) + + p1 = Point('p1') + p1.set_pos(pO, position) + p2 = Point('p2') + p2.set_pos(pO, position) + + with pytest.raises(ValueError): + _ = cylinder.geodesic_end_vectors(p1, p2) diff --git a/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6b8a7d3be358ab28e2438821ead7330fe7891611 Binary files /dev/null and b/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/anticommutator.cpython-312.pyc b/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/anticommutator.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bab93c5c4a2c0db4daad90d247124a4b4b37ee9b Binary files /dev/null and b/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/anticommutator.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/boson.cpython-312.pyc b/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/boson.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..44a58ed2f076dec14b413cae5fffe725d05f872c Binary files /dev/null and b/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/boson.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/cartesian.cpython-312.pyc b/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/cartesian.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..90ac68bcf22409ec320f25ac770b602f07c1fcd6 Binary files /dev/null and b/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/cartesian.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/cg.cpython-312.pyc b/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/cg.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e498735386130ebe0f38ce9061f565a65d2db2fc Binary files /dev/null and b/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/cg.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/circuitplot.cpython-312.pyc b/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/circuitplot.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8624f2203a5e375577b186b8d01bd5d7b7c73136 Binary files /dev/null and b/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/circuitplot.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/circuitutils.cpython-312.pyc b/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/circuitutils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c50247555a51f605331fee3e6c3687c01c088273 Binary files /dev/null and b/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/circuitutils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/commutator.cpython-312.pyc b/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/commutator.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..df91dcfe2b8c71133ffb7483d0afb9b294624881 Binary files /dev/null and b/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/commutator.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/constants.cpython-312.pyc b/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/constants.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cfaa42f736d544be3dacbef884858064e4073b24 Binary files /dev/null and b/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/constants.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/dagger.cpython-312.pyc b/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/dagger.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9790191209b39e58c5cf04c7d6e18e46ba3b5bf2 Binary files /dev/null and b/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/dagger.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/density.cpython-312.pyc b/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/density.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fa9d225373758350b971fb450303388c93e6c707 Binary files /dev/null and b/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/density.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/fermion.cpython-312.pyc b/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/fermion.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b68e878a56b617f1e2930de4e567480b95f52c52 Binary files /dev/null and b/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/fermion.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/gate.cpython-312.pyc b/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/gate.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2bf34892f730eaa1df007f431bb3356fe9faf86a Binary files /dev/null and b/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/gate.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/grover.cpython-312.pyc b/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/grover.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..379ae2357d40f2a00fbcf2f0761d220e3044ab35 Binary files /dev/null and b/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/grover.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/hilbert.cpython-312.pyc b/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/hilbert.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1171cca2d8f1b665d61bf846f6984248a602703d Binary files /dev/null and b/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/hilbert.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/identitysearch.cpython-312.pyc b/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/identitysearch.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..528eff7a2dd90a51b629def305670ec80f763a10 Binary files /dev/null and b/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/identitysearch.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/innerproduct.cpython-312.pyc b/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/innerproduct.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d49ddc076f9ee7526bd2e6bda361280c936400e5 Binary files /dev/null and b/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/innerproduct.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/matrixcache.cpython-312.pyc b/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/matrixcache.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..410e366880f589f7cfb546d0dd62ccdbce76d822 Binary files /dev/null and b/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/matrixcache.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/matrixutils.cpython-312.pyc b/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/matrixutils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3099fc34c69c8862fb1b3856f84c693b87d4233f Binary files /dev/null and b/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/matrixutils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/operator.cpython-312.pyc b/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/operator.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a552645c60333e342f3ad18d81e0c9cdbe0a9215 Binary files /dev/null and b/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/operator.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/operatorset.cpython-312.pyc b/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/operatorset.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..45d911ec3ee5523da6b9477c8854a5933b6f3bc2 Binary files /dev/null and b/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/operatorset.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/pauli.cpython-312.pyc b/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/pauli.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..feee19c840cb846f05d03c806566f4299c1c5af7 Binary files /dev/null and b/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/pauli.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/piab.cpython-312.pyc b/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/piab.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8aebf7efdc36ca28d275a56ccd4e2f825f1c8d1a Binary files /dev/null and b/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/piab.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/qapply.cpython-312.pyc b/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/qapply.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..00dc84d64e9eef32ecbc22ed70cf389fa9ff782a Binary files /dev/null and b/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/qapply.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/qasm.cpython-312.pyc b/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/qasm.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3511296c9c0e500b1da4aa97278f2c7bf7e25292 Binary files /dev/null and b/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/qasm.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/qexpr.cpython-312.pyc b/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/qexpr.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..59d7f80b9e50998a2da0179a9f34aecb656ff292 Binary files /dev/null and b/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/qexpr.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/qubit.cpython-312.pyc b/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/qubit.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..99be6692edf0dbeebd0a6e3256f5d176b8d987c9 Binary files /dev/null and b/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/qubit.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/represent.cpython-312.pyc b/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/represent.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..978ee4c6710b31704733fd512eca6f987de67ede Binary files /dev/null and b/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/represent.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/sho1d.cpython-312.pyc b/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/sho1d.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..baa237131c9fd31aa7d6d109d353613cc5eff64c Binary files /dev/null and b/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/sho1d.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/shor.cpython-312.pyc b/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/shor.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b6d21a393c2d05dad9063062cc50556dc356869d Binary files /dev/null and b/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/shor.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/state.cpython-312.pyc b/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/state.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..13e6a9a1f8d020d4805139c4b0c97a1176cd2721 Binary files /dev/null and b/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/state.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/tensorproduct.cpython-312.pyc b/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/tensorproduct.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bccc873f39ac56e0433839fc25a2d6e6a9c673c7 Binary files /dev/null and b/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/tensorproduct.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/trace.cpython-312.pyc b/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/trace.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..543703dbb3b71520f56bf7cb262d86c204358c6d Binary files /dev/null and b/lib/python3.12/site-packages/sympy/physics/quantum/__pycache__/trace.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sympy/stats/__init__.py b/lib/python3.12/site-packages/sympy/stats/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..adb79261954924305c1837555d7d47cd53b8430b --- /dev/null +++ b/lib/python3.12/site-packages/sympy/stats/__init__.py @@ -0,0 +1,202 @@ +""" +SymPy statistics module + +Introduces a random variable type into the SymPy language. + +Random variables may be declared using prebuilt functions such as +Normal, Exponential, Coin, Die, etc... or built with functions like FiniteRV. + +Queries on random expressions can be made using the functions + +========================= ============================= + Expression Meaning +------------------------- ----------------------------- + ``P(condition)`` Probability + ``E(expression)`` Expected value + ``H(expression)`` Entropy + ``variance(expression)`` Variance + ``density(expression)`` Probability Density Function + ``sample(expression)`` Produce a realization + ``where(condition)`` Where the condition is true +========================= ============================= + +Examples +======== + +>>> from sympy.stats import P, E, variance, Die, Normal +>>> from sympy import simplify +>>> X, Y = Die('X', 6), Die('Y', 6) # Define two six sided dice +>>> Z = Normal('Z', 0, 1) # Declare a Normal random variable with mean 0, std 1 +>>> P(X>3) # Probability X is greater than 3 +1/2 +>>> E(X+Y) # Expectation of the sum of two dice +7 +>>> variance(X+Y) # Variance of the sum of two dice +35/6 +>>> simplify(P(Z>1)) # Probability of Z being greater than 1 +1/2 - erf(sqrt(2)/2)/2 + + +One could also create custom distribution and define custom random variables +as follows: + +1. If you want to create a Continuous Random Variable: + +>>> from sympy.stats import ContinuousRV, P, E +>>> from sympy import exp, Symbol, Interval, oo +>>> x = Symbol('x') +>>> pdf = exp(-x) # pdf of the Continuous Distribution +>>> Z = ContinuousRV(x, pdf, set=Interval(0, oo)) +>>> E(Z) +1 +>>> P(Z > 5) +exp(-5) + +1.1 To create an instance of Continuous Distribution: + +>>> from sympy.stats import ContinuousDistributionHandmade +>>> from sympy import Lambda +>>> dist = ContinuousDistributionHandmade(Lambda(x, pdf), set=Interval(0, oo)) +>>> dist.pdf(x) +exp(-x) + +2. If you want to create a Discrete Random Variable: + +>>> from sympy.stats import DiscreteRV, P, E +>>> from sympy import Symbol, S +>>> p = S(1)/2 +>>> x = Symbol('x', integer=True, positive=True) +>>> pdf = p*(1 - p)**(x - 1) +>>> D = DiscreteRV(x, pdf, set=S.Naturals) +>>> E(D) +2 +>>> P(D > 3) +1/8 + +2.1 To create an instance of Discrete Distribution: + +>>> from sympy.stats import DiscreteDistributionHandmade +>>> from sympy import Lambda +>>> dist = DiscreteDistributionHandmade(Lambda(x, pdf), set=S.Naturals) +>>> dist.pdf(x) +2**(1 - x)/2 + +3. If you want to create a Finite Random Variable: + +>>> from sympy.stats import FiniteRV, P, E +>>> from sympy import Rational, Eq +>>> pmf = {1: Rational(1, 3), 2: Rational(1, 6), 3: Rational(1, 4), 4: Rational(1, 4)} +>>> X = FiniteRV('X', pmf) +>>> E(X) +29/12 +>>> P(X > 3) +1/4 + +3.1 To create an instance of Finite Distribution: + +>>> from sympy.stats import FiniteDistributionHandmade +>>> dist = FiniteDistributionHandmade(pmf) +>>> dist.pmf(x) +Lambda(x, Piecewise((1/3, Eq(x, 1)), (1/6, Eq(x, 2)), (1/4, Eq(x, 3) | Eq(x, 4)), (0, True))) +""" + +__all__ = [ + 'P', 'E', 'H', 'density', 'where', 'given', 'sample', 'cdf','median', + 'characteristic_function', 'pspace', 'sample_iter', 'variance', 'std', + 'skewness', 'kurtosis', 'covariance', 'dependent', 'entropy', 'independent', + 'random_symbols', 'correlation', 'factorial_moment', 'moment', 'cmoment', + 'sampling_density', 'moment_generating_function', 'smoment', 'quantile', + 'coskewness', 'sample_stochastic_process', + + 'FiniteRV', 'DiscreteUniform', 'Die', 'Bernoulli', 'Coin', 'Binomial', + 'BetaBinomial', 'Hypergeometric', 'Rademacher', 'IdealSoliton', 'RobustSoliton', + 'FiniteDistributionHandmade', + + 'ContinuousRV', 'Arcsin', 'Benini', 'Beta', 'BetaNoncentral', 'BetaPrime', + 'BoundedPareto', 'Cauchy', 'Chi', 'ChiNoncentral', 'ChiSquared', 'Dagum', 'Davis', 'Erlang', + 'ExGaussian', 'Exponential', 'ExponentialPower', 'FDistribution', + 'FisherZ', 'Frechet', 'Gamma', 'GammaInverse', 'Gompertz', 'Gumbel', + 'Kumaraswamy', 'Laplace', 'Levy', 'Logistic','LogCauchy', 'LogLogistic', 'LogitNormal', 'LogNormal', 'Lomax', + 'Moyal', 'Maxwell', 'Nakagami', 'Normal', 'GaussianInverse', 'Pareto', 'PowerFunction', + 'QuadraticU', 'RaisedCosine', 'Rayleigh','Reciprocal', 'StudentT', 'ShiftedGompertz', + 'Trapezoidal', 'Triangular', 'Uniform', 'UniformSum', 'VonMises', 'Wald', + 'Weibull', 'WignerSemicircle', 'ContinuousDistributionHandmade', + + 'FlorySchulz', 'Geometric','Hermite', 'Logarithmic', 'NegativeBinomial', 'Poisson', 'Skellam', + 'YuleSimon', 'Zeta', 'DiscreteRV', 'DiscreteDistributionHandmade', + + 'JointRV', 'Dirichlet', 'GeneralizedMultivariateLogGamma', + 'GeneralizedMultivariateLogGammaOmega', 'Multinomial', 'MultivariateBeta', + 'MultivariateEwens', 'MultivariateT', 'NegativeMultinomial', + 'NormalGamma', 'MultivariateNormal', 'MultivariateLaplace', 'marginal_distribution', + + 'StochasticProcess', 'DiscreteTimeStochasticProcess', + 'DiscreteMarkovChain', 'TransitionMatrixOf', 'StochasticStateSpaceOf', + 'GeneratorMatrixOf', 'ContinuousMarkovChain', 'BernoulliProcess', + 'PoissonProcess', 'WienerProcess', 'GammaProcess', + + 'CircularEnsemble', 'CircularUnitaryEnsemble', + 'CircularOrthogonalEnsemble', 'CircularSymplecticEnsemble', + 'GaussianEnsemble', 'GaussianUnitaryEnsemble', + 'GaussianOrthogonalEnsemble', 'GaussianSymplecticEnsemble', + 'joint_eigen_distribution', 'JointEigenDistribution', + 'level_spacing_distribution', + + 'MatrixGamma', 'Wishart', 'MatrixNormal', 'MatrixStudentT', + + 'Probability', 'Expectation', 'Variance', 'Covariance', 'Moment', + 'CentralMoment', + + 'ExpectationMatrix', 'VarianceMatrix', 'CrossCovarianceMatrix' + +] +from .rv_interface import (P, E, H, density, where, given, sample, cdf, median, + characteristic_function, pspace, sample_iter, variance, std, skewness, + kurtosis, covariance, dependent, entropy, independent, random_symbols, + correlation, factorial_moment, moment, cmoment, sampling_density, + moment_generating_function, smoment, quantile, coskewness, + sample_stochastic_process) + +from .frv_types import (FiniteRV, DiscreteUniform, Die, Bernoulli, Coin, + Binomial, BetaBinomial, Hypergeometric, Rademacher, + FiniteDistributionHandmade, IdealSoliton, RobustSoliton) + +from .crv_types import (ContinuousRV, Arcsin, Benini, Beta, BetaNoncentral, + BetaPrime, BoundedPareto, Cauchy, Chi, ChiNoncentral, ChiSquared, + Dagum, Davis, Erlang, ExGaussian, Exponential, ExponentialPower, + FDistribution, FisherZ, Frechet, Gamma, GammaInverse, GaussianInverse, + Gompertz, Gumbel, Kumaraswamy, Laplace, Levy, Logistic, LogCauchy, + LogLogistic, LogitNormal, LogNormal, Lomax, Maxwell, Moyal, Nakagami, + Normal, Pareto, QuadraticU, RaisedCosine, Rayleigh, Reciprocal, + StudentT, PowerFunction, ShiftedGompertz, Trapezoidal, Triangular, + Uniform, UniformSum, VonMises, Wald, Weibull, WignerSemicircle, + ContinuousDistributionHandmade) + +from .drv_types import (FlorySchulz, Geometric, Hermite, Logarithmic, NegativeBinomial, Poisson, + Skellam, YuleSimon, Zeta, DiscreteRV, DiscreteDistributionHandmade) + +from .joint_rv_types import (JointRV, Dirichlet, + GeneralizedMultivariateLogGamma, GeneralizedMultivariateLogGammaOmega, + Multinomial, MultivariateBeta, MultivariateEwens, MultivariateT, + NegativeMultinomial, NormalGamma, MultivariateNormal, MultivariateLaplace, + marginal_distribution) + +from .stochastic_process_types import (StochasticProcess, + DiscreteTimeStochasticProcess, DiscreteMarkovChain, + TransitionMatrixOf, StochasticStateSpaceOf, GeneratorMatrixOf, + ContinuousMarkovChain, BernoulliProcess, PoissonProcess, WienerProcess, + GammaProcess) + +from .random_matrix_models import (CircularEnsemble, CircularUnitaryEnsemble, + CircularOrthogonalEnsemble, CircularSymplecticEnsemble, + GaussianEnsemble, GaussianUnitaryEnsemble, GaussianOrthogonalEnsemble, + GaussianSymplecticEnsemble, joint_eigen_distribution, + JointEigenDistribution, level_spacing_distribution) + +from .matrix_distributions import MatrixGamma, Wishart, MatrixNormal, MatrixStudentT + +from .symbolic_probability import (Probability, Expectation, Variance, + Covariance, Moment, CentralMoment) + +from .symbolic_multivariate_probability import (ExpectationMatrix, VarianceMatrix, + CrossCovarianceMatrix) diff --git a/lib/python3.12/site-packages/sympy/stats/__pycache__/compound_rv.cpython-312.pyc b/lib/python3.12/site-packages/sympy/stats/__pycache__/compound_rv.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..631062096a2ee03604254b833444005b16fa1f29 Binary files /dev/null and b/lib/python3.12/site-packages/sympy/stats/__pycache__/compound_rv.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sympy/stats/__pycache__/drv.cpython-312.pyc b/lib/python3.12/site-packages/sympy/stats/__pycache__/drv.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..84012e2945886a9146bf35960020f634f9c9843b Binary files /dev/null and b/lib/python3.12/site-packages/sympy/stats/__pycache__/drv.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sympy/stats/__pycache__/drv_types.cpython-312.pyc b/lib/python3.12/site-packages/sympy/stats/__pycache__/drv_types.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..21c0b6d47053b25ff829c5c9e81802e8c0ce314f Binary files /dev/null and b/lib/python3.12/site-packages/sympy/stats/__pycache__/drv_types.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sympy/stats/__pycache__/frv.cpython-312.pyc b/lib/python3.12/site-packages/sympy/stats/__pycache__/frv.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..77fd0bb44e933c9b38b41e2e7c631ea01bf86c25 Binary files /dev/null and b/lib/python3.12/site-packages/sympy/stats/__pycache__/frv.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sympy/stats/__pycache__/joint_rv.cpython-312.pyc b/lib/python3.12/site-packages/sympy/stats/__pycache__/joint_rv.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6c35914c1e71a5edefa58cb10ec19f0c0333c0b7 Binary files /dev/null and b/lib/python3.12/site-packages/sympy/stats/__pycache__/joint_rv.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sympy/stats/__pycache__/joint_rv_types.cpython-312.pyc b/lib/python3.12/site-packages/sympy/stats/__pycache__/joint_rv_types.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a74a5ba1c069b54d426366e6d86f09b13e948cf4 Binary files /dev/null and b/lib/python3.12/site-packages/sympy/stats/__pycache__/joint_rv_types.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sympy/stats/__pycache__/rv.cpython-312.pyc b/lib/python3.12/site-packages/sympy/stats/__pycache__/rv.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d5733e525d46c84bfe3563508a4304838c7d7ef9 Binary files /dev/null and b/lib/python3.12/site-packages/sympy/stats/__pycache__/rv.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sympy/stats/__pycache__/rv_interface.cpython-312.pyc b/lib/python3.12/site-packages/sympy/stats/__pycache__/rv_interface.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..70f6f9fb49062f0a14d4600d1929e13d5bb69f13 Binary files /dev/null and b/lib/python3.12/site-packages/sympy/stats/__pycache__/rv_interface.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sympy/stats/__pycache__/stochastic_process.cpython-312.pyc b/lib/python3.12/site-packages/sympy/stats/__pycache__/stochastic_process.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3c4510fe60a020e37b0a5d9b3cd4f858e2cee282 Binary files /dev/null and b/lib/python3.12/site-packages/sympy/stats/__pycache__/stochastic_process.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sympy/stats/__pycache__/symbolic_probability.cpython-312.pyc b/lib/python3.12/site-packages/sympy/stats/__pycache__/symbolic_probability.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2533554a37086b6ee92eeeba9f5e8b9e463298fd Binary files /dev/null and b/lib/python3.12/site-packages/sympy/stats/__pycache__/symbolic_probability.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sympy/stats/compound_rv.py b/lib/python3.12/site-packages/sympy/stats/compound_rv.py new file mode 100644 index 0000000000000000000000000000000000000000..27555f4233fe691bac303800a87736205acbdee6 --- /dev/null +++ b/lib/python3.12/site-packages/sympy/stats/compound_rv.py @@ -0,0 +1,223 @@ +from sympy.concrete.summations import Sum +from sympy.core.basic import Basic +from sympy.core.function import Lambda +from sympy.core.symbol import Dummy +from sympy.integrals.integrals import Integral +from sympy.stats.rv import (NamedArgsMixin, random_symbols, _symbol_converter, + PSpace, RandomSymbol, is_random, Distribution) +from sympy.stats.crv import ContinuousDistribution, SingleContinuousPSpace +from sympy.stats.drv import DiscreteDistribution, SingleDiscretePSpace +from sympy.stats.frv import SingleFiniteDistribution, SingleFinitePSpace +from sympy.stats.crv_types import ContinuousDistributionHandmade +from sympy.stats.drv_types import DiscreteDistributionHandmade +from sympy.stats.frv_types import FiniteDistributionHandmade + + +class CompoundPSpace(PSpace): + """ + A temporary Probability Space for the Compound Distribution. After + Marginalization, this returns the corresponding Probability Space of the + parent distribution. + """ + + def __new__(cls, s, distribution): + s = _symbol_converter(s) + if isinstance(distribution, ContinuousDistribution): + return SingleContinuousPSpace(s, distribution) + if isinstance(distribution, DiscreteDistribution): + return SingleDiscretePSpace(s, distribution) + if isinstance(distribution, SingleFiniteDistribution): + return SingleFinitePSpace(s, distribution) + if not isinstance(distribution, CompoundDistribution): + raise ValueError("%s should be an isinstance of " + "CompoundDistribution"%(distribution)) + return Basic.__new__(cls, s, distribution) + + @property + def value(self): + return RandomSymbol(self.symbol, self) + + @property + def symbol(self): + return self.args[0] + + @property + def is_Continuous(self): + return self.distribution.is_Continuous + + @property + def is_Finite(self): + return self.distribution.is_Finite + + @property + def is_Discrete(self): + return self.distribution.is_Discrete + + @property + def distribution(self): + return self.args[1] + + @property + def pdf(self): + return self.distribution.pdf(self.symbol) + + @property + def set(self): + return self.distribution.set + + @property + def domain(self): + return self._get_newpspace().domain + + def _get_newpspace(self, evaluate=False): + x = Dummy('x') + parent_dist = self.distribution.args[0] + func = Lambda(x, self.distribution.pdf(x, evaluate)) + new_pspace = self._transform_pspace(self.symbol, parent_dist, func) + if new_pspace is not None: + return new_pspace + message = ("Compound Distribution for %s is not implemented yet" % str(parent_dist)) + raise NotImplementedError(message) + + def _transform_pspace(self, sym, dist, pdf): + """ + This function returns the new pspace of the distribution using handmade + Distributions and their corresponding pspace. + """ + pdf = Lambda(sym, pdf(sym)) + _set = dist.set + if isinstance(dist, ContinuousDistribution): + return SingleContinuousPSpace(sym, ContinuousDistributionHandmade(pdf, _set)) + elif isinstance(dist, DiscreteDistribution): + return SingleDiscretePSpace(sym, DiscreteDistributionHandmade(pdf, _set)) + elif isinstance(dist, SingleFiniteDistribution): + dens = {k: pdf(k) for k in _set} + return SingleFinitePSpace(sym, FiniteDistributionHandmade(dens)) + + def compute_density(self, expr, *, compound_evaluate=True, **kwargs): + new_pspace = self._get_newpspace(compound_evaluate) + expr = expr.subs({self.value: new_pspace.value}) + return new_pspace.compute_density(expr, **kwargs) + + def compute_cdf(self, expr, *, compound_evaluate=True, **kwargs): + new_pspace = self._get_newpspace(compound_evaluate) + expr = expr.subs({self.value: new_pspace.value}) + return new_pspace.compute_cdf(expr, **kwargs) + + def compute_expectation(self, expr, rvs=None, evaluate=False, **kwargs): + new_pspace = self._get_newpspace(evaluate) + expr = expr.subs({self.value: new_pspace.value}) + if rvs: + rvs = rvs.subs({self.value: new_pspace.value}) + if isinstance(new_pspace, SingleFinitePSpace): + return new_pspace.compute_expectation(expr, rvs, **kwargs) + return new_pspace.compute_expectation(expr, rvs, evaluate, **kwargs) + + def probability(self, condition, *, compound_evaluate=True, **kwargs): + new_pspace = self._get_newpspace(compound_evaluate) + condition = condition.subs({self.value: new_pspace.value}) + return new_pspace.probability(condition) + + def conditional_space(self, condition, *, compound_evaluate=True, **kwargs): + new_pspace = self._get_newpspace(compound_evaluate) + condition = condition.subs({self.value: new_pspace.value}) + return new_pspace.conditional_space(condition) + + +class CompoundDistribution(Distribution, NamedArgsMixin): + """ + Class for Compound Distributions. + + Parameters + ========== + + dist : Distribution + Distribution must contain a random parameter + + Examples + ======== + + >>> from sympy.stats.compound_rv import CompoundDistribution + >>> from sympy.stats.crv_types import NormalDistribution + >>> from sympy.stats import Normal + >>> from sympy.abc import x + >>> X = Normal('X', 2, 4) + >>> N = NormalDistribution(X, 4) + >>> C = CompoundDistribution(N) + >>> C.set + Interval(-oo, oo) + >>> C.pdf(x, evaluate=True).simplify() + exp(-x**2/64 + x/16 - 1/16)/(8*sqrt(pi)) + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Compound_probability_distribution + + """ + + def __new__(cls, dist): + if not isinstance(dist, (ContinuousDistribution, + SingleFiniteDistribution, DiscreteDistribution)): + message = "Compound Distribution for %s is not implemented yet" % str(dist) + raise NotImplementedError(message) + if not cls._compound_check(dist): + return dist + return Basic.__new__(cls, dist) + + @property + def set(self): + return self.args[0].set + + @property + def is_Continuous(self): + return isinstance(self.args[0], ContinuousDistribution) + + @property + def is_Finite(self): + return isinstance(self.args[0], SingleFiniteDistribution) + + @property + def is_Discrete(self): + return isinstance(self.args[0], DiscreteDistribution) + + def pdf(self, x, evaluate=False): + dist = self.args[0] + randoms = [rv for rv in dist.args if is_random(rv)] + if isinstance(dist, SingleFiniteDistribution): + y = Dummy('y', integer=True, negative=False) + expr = dist.pmf(y) + else: + y = Dummy('y') + expr = dist.pdf(y) + for rv in randoms: + expr = self._marginalise(expr, rv, evaluate) + return Lambda(y, expr)(x) + + def _marginalise(self, expr, rv, evaluate): + if isinstance(rv.pspace.distribution, SingleFiniteDistribution): + rv_dens = rv.pspace.distribution.pmf(rv) + else: + rv_dens = rv.pspace.distribution.pdf(rv) + rv_dom = rv.pspace.domain.set + if rv.pspace.is_Discrete or rv.pspace.is_Finite: + expr = Sum(expr*rv_dens, (rv, rv_dom._inf, + rv_dom._sup)) + else: + expr = Integral(expr*rv_dens, (rv, rv_dom._inf, + rv_dom._sup)) + if evaluate: + return expr.doit() + return expr + + @classmethod + def _compound_check(self, dist): + """ + Checks if the given distribution contains random parameters. + """ + randoms = [] + for arg in dist.args: + randoms.extend(random_symbols(arg)) + if len(randoms) == 0: + return False + return True diff --git a/lib/python3.12/site-packages/sympy/stats/crv.py b/lib/python3.12/site-packages/sympy/stats/crv.py new file mode 100644 index 0000000000000000000000000000000000000000..0a5184029679f663c83d81aa6c1b6ca4d948c70f --- /dev/null +++ b/lib/python3.12/site-packages/sympy/stats/crv.py @@ -0,0 +1,570 @@ +""" +Continuous Random Variables Module + +See Also +======== +sympy.stats.crv_types +sympy.stats.rv +sympy.stats.frv +""" + + +from sympy.core.basic import Basic +from sympy.core.cache import cacheit +from sympy.core.function import Lambda, PoleError +from sympy.core.numbers import (I, nan, oo) +from sympy.core.relational import (Eq, Ne) +from sympy.core.singleton import S +from sympy.core.symbol import (Dummy, symbols) +from sympy.core.sympify import _sympify, sympify +from sympy.functions.combinatorial.factorials import factorial +from sympy.functions.elementary.exponential import exp +from sympy.functions.elementary.piecewise import Piecewise +from sympy.functions.special.delta_functions import DiracDelta +from sympy.integrals.integrals import (Integral, integrate) +from sympy.logic.boolalg import (And, Or) +from sympy.polys.polyerrors import PolynomialError +from sympy.polys.polytools import poly +from sympy.series.series import series +from sympy.sets.sets import (FiniteSet, Intersection, Interval, Union) +from sympy.solvers.solveset import solveset +from sympy.solvers.inequalities import reduce_rational_inequalities +from sympy.stats.rv import (RandomDomain, SingleDomain, ConditionalDomain, is_random, + ProductDomain, PSpace, SinglePSpace, random_symbols, NamedArgsMixin, Distribution) + + +class ContinuousDomain(RandomDomain): + """ + A domain with continuous support + + Represented using symbols and Intervals. + """ + is_Continuous = True + + def as_boolean(self): + raise NotImplementedError("Not Implemented for generic Domains") + + +class SingleContinuousDomain(ContinuousDomain, SingleDomain): + """ + A univariate domain with continuous support + + Represented using a single symbol and interval. + """ + def compute_expectation(self, expr, variables=None, **kwargs): + if variables is None: + variables = self.symbols + if not variables: + return expr + if frozenset(variables) != frozenset(self.symbols): + raise ValueError("Values should be equal") + # assumes only intervals + return Integral(expr, (self.symbol, self.set), **kwargs) + + def as_boolean(self): + return self.set.as_relational(self.symbol) + + +class ProductContinuousDomain(ProductDomain, ContinuousDomain): + """ + A collection of independent domains with continuous support + """ + + def compute_expectation(self, expr, variables=None, **kwargs): + if variables is None: + variables = self.symbols + for domain in self.domains: + domain_vars = frozenset(variables) & frozenset(domain.symbols) + if domain_vars: + expr = domain.compute_expectation(expr, domain_vars, **kwargs) + return expr + + def as_boolean(self): + return And(*[domain.as_boolean() for domain in self.domains]) + + +class ConditionalContinuousDomain(ContinuousDomain, ConditionalDomain): + """ + A domain with continuous support that has been further restricted by a + condition such as $x > 3$. + """ + + def compute_expectation(self, expr, variables=None, **kwargs): + if variables is None: + variables = self.symbols + if not variables: + return expr + # Extract the full integral + fullintgrl = self.fulldomain.compute_expectation(expr, variables) + # separate into integrand and limits + integrand, limits = fullintgrl.function, list(fullintgrl.limits) + + conditions = [self.condition] + while conditions: + cond = conditions.pop() + if cond.is_Boolean: + if isinstance(cond, And): + conditions.extend(cond.args) + elif isinstance(cond, Or): + raise NotImplementedError("Or not implemented here") + elif cond.is_Relational: + if cond.is_Equality: + # Add the appropriate Delta to the integrand + integrand *= DiracDelta(cond.lhs - cond.rhs) + else: + symbols = cond.free_symbols & set(self.symbols) + if len(symbols) != 1: # Can't handle x > y + raise NotImplementedError( + "Multivariate Inequalities not yet implemented") + # Can handle x > 0 + symbol = symbols.pop() + # Find the limit with x, such as (x, -oo, oo) + for i, limit in enumerate(limits): + if limit[0] == symbol: + # Make condition into an Interval like [0, oo] + cintvl = reduce_rational_inequalities_wrap( + cond, symbol) + # Make limit into an Interval like [-oo, oo] + lintvl = Interval(limit[1], limit[2]) + # Intersect them to get [0, oo] + intvl = cintvl.intersect(lintvl) + # Put back into limits list + limits[i] = (symbol, intvl.left, intvl.right) + else: + raise TypeError( + "Condition %s is not a relational or Boolean" % cond) + + return Integral(integrand, *limits, **kwargs) + + def as_boolean(self): + return And(self.fulldomain.as_boolean(), self.condition) + + @property + def set(self): + if len(self.symbols) == 1: + return (self.fulldomain.set & reduce_rational_inequalities_wrap( + self.condition, tuple(self.symbols)[0])) + else: + raise NotImplementedError( + "Set of Conditional Domain not Implemented") + + +class ContinuousDistribution(Distribution): + def __call__(self, *args): + return self.pdf(*args) + + +class SingleContinuousDistribution(ContinuousDistribution, NamedArgsMixin): + """ Continuous distribution of a single variable. + + Explanation + =========== + + Serves as superclass for Normal/Exponential/UniformDistribution etc.... + + Represented by parameters for each of the specific classes. E.g + NormalDistribution is represented by a mean and standard deviation. + + Provides methods for pdf, cdf, and sampling. + + See Also + ======== + + sympy.stats.crv_types.* + """ + + set = Interval(-oo, oo) + + def __new__(cls, *args): + args = list(map(sympify, args)) + return Basic.__new__(cls, *args) + + @staticmethod + def check(*args): + pass + + @cacheit + def compute_cdf(self, **kwargs): + """ Compute the CDF from the PDF. + + Returns a Lambda. + """ + x, z = symbols('x, z', real=True, cls=Dummy) + left_bound = self.set.start + + # CDF is integral of PDF from left bound to z + pdf = self.pdf(x) + cdf = integrate(pdf.doit(), (x, left_bound, z), **kwargs) + # CDF Ensure that CDF left of left_bound is zero + cdf = Piecewise((cdf, z >= left_bound), (0, True)) + return Lambda(z, cdf) + + def _cdf(self, x): + return None + + def cdf(self, x, **kwargs): + """ Cumulative density function """ + if len(kwargs) == 0: + cdf = self._cdf(x) + if cdf is not None: + return cdf + return self.compute_cdf(**kwargs)(x) + + @cacheit + def compute_characteristic_function(self, **kwargs): + """ Compute the characteristic function from the PDF. + + Returns a Lambda. + """ + x, t = symbols('x, t', real=True, cls=Dummy) + pdf = self.pdf(x) + cf = integrate(exp(I*t*x)*pdf, (x, self.set)) + return Lambda(t, cf) + + def _characteristic_function(self, t): + return None + + def characteristic_function(self, t, **kwargs): + """ Characteristic function """ + if len(kwargs) == 0: + cf = self._characteristic_function(t) + if cf is not None: + return cf + return self.compute_characteristic_function(**kwargs)(t) + + @cacheit + def compute_moment_generating_function(self, **kwargs): + """ Compute the moment generating function from the PDF. + + Returns a Lambda. + """ + x, t = symbols('x, t', real=True, cls=Dummy) + pdf = self.pdf(x) + mgf = integrate(exp(t * x) * pdf, (x, self.set)) + return Lambda(t, mgf) + + def _moment_generating_function(self, t): + return None + + def moment_generating_function(self, t, **kwargs): + """ Moment generating function """ + if not kwargs: + mgf = self._moment_generating_function(t) + if mgf is not None: + return mgf + return self.compute_moment_generating_function(**kwargs)(t) + + def expectation(self, expr, var, evaluate=True, **kwargs): + """ Expectation of expression over distribution """ + if evaluate: + try: + p = poly(expr, var) + if p.is_zero: + return S.Zero + t = Dummy('t', real=True) + mgf = self._moment_generating_function(t) + if mgf is None: + return integrate(expr * self.pdf(var), (var, self.set), **kwargs) + deg = p.degree() + taylor = poly(series(mgf, t, 0, deg + 1).removeO(), t) + result = 0 + for k in range(deg+1): + result += p.coeff_monomial(var ** k) * taylor.coeff_monomial(t ** k) * factorial(k) + return result + except PolynomialError: + return integrate(expr * self.pdf(var), (var, self.set), **kwargs) + else: + return Integral(expr * self.pdf(var), (var, self.set), **kwargs) + + @cacheit + def compute_quantile(self, **kwargs): + """ Compute the Quantile from the PDF. + + Returns a Lambda. + """ + x, p = symbols('x, p', real=True, cls=Dummy) + left_bound = self.set.start + + pdf = self.pdf(x) + cdf = integrate(pdf, (x, left_bound, x), **kwargs) + quantile = solveset(cdf - p, x, self.set) + return Lambda(p, Piecewise((quantile, (p >= 0) & (p <= 1) ), (nan, True))) + + def _quantile(self, x): + return None + + def quantile(self, x, **kwargs): + """ Cumulative density function """ + if len(kwargs) == 0: + quantile = self._quantile(x) + if quantile is not None: + return quantile + return self.compute_quantile(**kwargs)(x) + + +class ContinuousPSpace(PSpace): + """ Continuous Probability Space + + Represents the likelihood of an event space defined over a continuum. + + Represented with a ContinuousDomain and a PDF (Lambda-Like) + """ + + is_Continuous = True + is_real = True + + @property + def pdf(self): + return self.density(*self.domain.symbols) + + def compute_expectation(self, expr, rvs=None, evaluate=False, **kwargs): + if rvs is None: + rvs = self.values + else: + rvs = frozenset(rvs) + + expr = expr.xreplace({rv: rv.symbol for rv in rvs}) + + domain_symbols = frozenset(rv.symbol for rv in rvs) + + return self.domain.compute_expectation(self.pdf * expr, + domain_symbols, **kwargs) + + def compute_density(self, expr, **kwargs): + # Common case Density(X) where X in self.values + if expr in self.values: + # Marginalize all other random symbols out of the density + randomsymbols = tuple(set(self.values) - frozenset([expr])) + symbols = tuple(rs.symbol for rs in randomsymbols) + pdf = self.domain.compute_expectation(self.pdf, symbols, **kwargs) + return Lambda(expr.symbol, pdf) + + z = Dummy('z', real=True) + return Lambda(z, self.compute_expectation(DiracDelta(expr - z), **kwargs)) + + @cacheit + def compute_cdf(self, expr, **kwargs): + if not self.domain.set.is_Interval: + raise ValueError( + "CDF not well defined on multivariate expressions") + + d = self.compute_density(expr, **kwargs) + x, z = symbols('x, z', real=True, cls=Dummy) + left_bound = self.domain.set.start + + # CDF is integral of PDF from left bound to z + cdf = integrate(d(x), (x, left_bound, z), **kwargs) + # CDF Ensure that CDF left of left_bound is zero + cdf = Piecewise((cdf, z >= left_bound), (0, True)) + return Lambda(z, cdf) + + @cacheit + def compute_characteristic_function(self, expr, **kwargs): + if not self.domain.set.is_Interval: + raise NotImplementedError("Characteristic function of multivariate expressions not implemented") + + d = self.compute_density(expr, **kwargs) + x, t = symbols('x, t', real=True, cls=Dummy) + cf = integrate(exp(I*t*x)*d(x), (x, -oo, oo), **kwargs) + return Lambda(t, cf) + + @cacheit + def compute_moment_generating_function(self, expr, **kwargs): + if not self.domain.set.is_Interval: + raise NotImplementedError("Moment generating function of multivariate expressions not implemented") + + d = self.compute_density(expr, **kwargs) + x, t = symbols('x, t', real=True, cls=Dummy) + mgf = integrate(exp(t * x) * d(x), (x, -oo, oo), **kwargs) + return Lambda(t, mgf) + + @cacheit + def compute_quantile(self, expr, **kwargs): + if not self.domain.set.is_Interval: + raise ValueError( + "Quantile not well defined on multivariate expressions") + + d = self.compute_cdf(expr, **kwargs) + x = Dummy('x', real=True) + p = Dummy('p', positive=True) + + quantile = solveset(d(x) - p, x, self.set) + + return Lambda(p, quantile) + + def probability(self, condition, **kwargs): + z = Dummy('z', real=True) + cond_inv = False + if isinstance(condition, Ne): + condition = Eq(condition.args[0], condition.args[1]) + cond_inv = True + # Univariate case can be handled by where + try: + domain = self.where(condition) + rv = [rv for rv in self.values if rv.symbol == domain.symbol][0] + # Integrate out all other random variables + pdf = self.compute_density(rv, **kwargs) + # return S.Zero if `domain` is empty set + if domain.set is S.EmptySet or isinstance(domain.set, FiniteSet): + return S.Zero if not cond_inv else S.One + if isinstance(domain.set, Union): + return sum( + Integral(pdf(z), (z, subset), **kwargs) for subset in + domain.set.args if isinstance(subset, Interval)) + # Integrate out the last variable over the special domain + return Integral(pdf(z), (z, domain.set), **kwargs) + + # Other cases can be turned into univariate case + # by computing a density handled by density computation + except NotImplementedError: + from sympy.stats.rv import density + expr = condition.lhs - condition.rhs + if not is_random(expr): + dens = self.density + comp = condition.rhs + else: + dens = density(expr, **kwargs) + comp = 0 + if not isinstance(dens, ContinuousDistribution): + from sympy.stats.crv_types import ContinuousDistributionHandmade + dens = ContinuousDistributionHandmade(dens, set=self.domain.set) + # Turn problem into univariate case + space = SingleContinuousPSpace(z, dens) + result = space.probability(condition.__class__(space.value, comp)) + return result if not cond_inv else S.One - result + + def where(self, condition): + rvs = frozenset(random_symbols(condition)) + if not (len(rvs) == 1 and rvs.issubset(self.values)): + raise NotImplementedError( + "Multiple continuous random variables not supported") + rv = tuple(rvs)[0] + interval = reduce_rational_inequalities_wrap(condition, rv) + interval = interval.intersect(self.domain.set) + return SingleContinuousDomain(rv.symbol, interval) + + def conditional_space(self, condition, normalize=True, **kwargs): + condition = condition.xreplace({rv: rv.symbol for rv in self.values}) + domain = ConditionalContinuousDomain(self.domain, condition) + if normalize: + # create a clone of the variable to + # make sure that variables in nested integrals are different + # from the variables outside the integral + # this makes sure that they are evaluated separately + # and in the correct order + replacement = {rv: Dummy(str(rv)) for rv in self.symbols} + norm = domain.compute_expectation(self.pdf, **kwargs) + pdf = self.pdf / norm.xreplace(replacement) + # XXX: Converting set to tuple. The order matters to Lambda though + # so we shouldn't be starting with a set here... + density = Lambda(tuple(domain.symbols), pdf) + + return ContinuousPSpace(domain, density) + + +class SingleContinuousPSpace(ContinuousPSpace, SinglePSpace): + """ + A continuous probability space over a single univariate variable. + + These consist of a Symbol and a SingleContinuousDistribution + + This class is normally accessed through the various random variable + functions, Normal, Exponential, Uniform, etc.... + """ + + @property + def set(self): + return self.distribution.set + + @property + def domain(self): + return SingleContinuousDomain(sympify(self.symbol), self.set) + + def sample(self, size=(), library='scipy', seed=None): + """ + Internal sample method. + + Returns dictionary mapping RandomSymbol to realization value. + """ + return {self.value: self.distribution.sample(size, library=library, seed=seed)} + + def compute_expectation(self, expr, rvs=None, evaluate=False, **kwargs): + rvs = rvs or (self.value,) + if self.value not in rvs: + return expr + + expr = _sympify(expr) + expr = expr.xreplace({rv: rv.symbol for rv in rvs}) + + x = self.value.symbol + try: + return self.distribution.expectation(expr, x, evaluate=evaluate, **kwargs) + except PoleError: + return Integral(expr * self.pdf, (x, self.set), **kwargs) + + def compute_cdf(self, expr, **kwargs): + if expr == self.value: + z = Dummy("z", real=True) + return Lambda(z, self.distribution.cdf(z, **kwargs)) + else: + return ContinuousPSpace.compute_cdf(self, expr, **kwargs) + + def compute_characteristic_function(self, expr, **kwargs): + if expr == self.value: + t = Dummy("t", real=True) + return Lambda(t, self.distribution.characteristic_function(t, **kwargs)) + else: + return ContinuousPSpace.compute_characteristic_function(self, expr, **kwargs) + + def compute_moment_generating_function(self, expr, **kwargs): + if expr == self.value: + t = Dummy("t", real=True) + return Lambda(t, self.distribution.moment_generating_function(t, **kwargs)) + else: + return ContinuousPSpace.compute_moment_generating_function(self, expr, **kwargs) + + def compute_density(self, expr, **kwargs): + # https://en.wikipedia.org/wiki/Random_variable#Functions_of_random_variables + if expr == self.value: + return self.density + y = Dummy('y', real=True) + + gs = solveset(expr - y, self.value, S.Reals) + + if isinstance(gs, Intersection): + if len(gs.args) == 2 and gs.args[0] is S.Reals: + gs = gs.args[1] + if not gs.is_FiniteSet: + raise ValueError("Can not solve %s for %s" % (expr, self.value)) + fx = self.compute_density(self.value) + fy = sum(fx(g) * abs(g.diff(y)) for g in gs) + return Lambda(y, fy) + + def compute_quantile(self, expr, **kwargs): + + if expr == self.value: + p = Dummy("p", real=True) + return Lambda(p, self.distribution.quantile(p, **kwargs)) + else: + return ContinuousPSpace.compute_quantile(self, expr, **kwargs) + +def _reduce_inequalities(conditions, var, **kwargs): + try: + return reduce_rational_inequalities(conditions, var, **kwargs) + except PolynomialError: + raise ValueError("Reduction of condition failed %s\n" % conditions[0]) + + +def reduce_rational_inequalities_wrap(condition, var): + if condition.is_Relational: + return _reduce_inequalities([[condition]], var, relational=False) + if isinstance(condition, Or): + return Union(*[_reduce_inequalities([[arg]], var, relational=False) + for arg in condition.args]) + if isinstance(condition, And): + intervals = [_reduce_inequalities([[arg]], var, relational=False) + for arg in condition.args] + I = intervals[0] + for i in intervals: + I = I.intersect(i) + return I diff --git a/lib/python3.12/site-packages/sympy/stats/crv_types.py b/lib/python3.12/site-packages/sympy/stats/crv_types.py new file mode 100644 index 0000000000000000000000000000000000000000..073e7350fdf80aac39ecd1dd607488a8b76187e3 --- /dev/null +++ b/lib/python3.12/site-packages/sympy/stats/crv_types.py @@ -0,0 +1,4732 @@ +""" +Continuous Random Variables - Prebuilt variables + +Contains +======== +Arcsin +Benini +Beta +BetaNoncentral +BetaPrime +BoundedPareto +Cauchy +Chi +ChiNoncentral +ChiSquared +Dagum +Davis +Erlang +ExGaussian +Exponential +ExponentialPower +FDistribution +FisherZ +Frechet +Gamma +GammaInverse +Gumbel +Gompertz +Kumaraswamy +Laplace +Levy +LogCauchy +Logistic +LogLogistic +LogitNormal +LogNormal +Lomax +Maxwell +Moyal +Nakagami +Normal +Pareto +PowerFunction +QuadraticU +RaisedCosine +Rayleigh +Reciprocal +ShiftedGompertz +StudentT +Trapezoidal +Triangular +Uniform +UniformSum +VonMises +Wald +Weibull +WignerSemicircle +""" + + + +from sympy.functions.elementary.exponential import exp +from sympy.functions.elementary.trigonometric import (atan, cos, sin, tan) +from sympy.functions.special.bessel import (besseli, besselj, besselk) +from sympy.functions.special.beta_functions import beta as beta_fn +from sympy.concrete.summations import Sum +from sympy.core.basic import Basic +from sympy.core.function import Lambda +from sympy.core.numbers import (I, Rational, pi) +from sympy.core.relational import (Eq, Ne) +from sympy.core.singleton import S +from sympy.core.symbol import Dummy +from sympy.core.sympify import sympify +from sympy.functions.combinatorial.factorials import (binomial, factorial) +from sympy.functions.elementary.complexes import (Abs, sign) +from sympy.functions.elementary.exponential import log +from sympy.functions.elementary.hyperbolic import sinh +from sympy.functions.elementary.integers import floor +from sympy.functions.elementary.miscellaneous import sqrt, Max, Min +from sympy.functions.elementary.piecewise import Piecewise +from sympy.functions.elementary.trigonometric import asin +from sympy.functions.special.error_functions import (erf, erfc, erfi, erfinv, expint) +from sympy.functions.special.gamma_functions import (gamma, lowergamma, uppergamma) +from sympy.functions.special.zeta_functions import zeta +from sympy.functions.special.hyper import hyper +from sympy.integrals.integrals import integrate +from sympy.logic.boolalg import And +from sympy.sets.sets import Interval +from sympy.matrices import MatrixBase +from sympy.stats.crv import SingleContinuousPSpace, SingleContinuousDistribution +from sympy.stats.rv import _value_check, is_random + +oo = S.Infinity + +__all__ = ['ContinuousRV', +'Arcsin', +'Benini', +'Beta', +'BetaNoncentral', +'BetaPrime', +'BoundedPareto', +'Cauchy', +'Chi', +'ChiNoncentral', +'ChiSquared', +'Dagum', +'Davis', +'Erlang', +'ExGaussian', +'Exponential', +'ExponentialPower', +'FDistribution', +'FisherZ', +'Frechet', +'Gamma', +'GammaInverse', +'Gompertz', +'Gumbel', +'Kumaraswamy', +'Laplace', +'Levy', +'LogCauchy', +'Logistic', +'LogLogistic', +'LogitNormal', +'LogNormal', +'Lomax', +'Maxwell', +'Moyal', +'Nakagami', +'Normal', +'GaussianInverse', +'Pareto', +'PowerFunction', +'QuadraticU', +'RaisedCosine', +'Rayleigh', +'Reciprocal', +'StudentT', +'ShiftedGompertz', +'Trapezoidal', +'Triangular', +'Uniform', +'UniformSum', +'VonMises', +'Wald', +'Weibull', +'WignerSemicircle', +] + + +@is_random.register(MatrixBase) +def _(x): + return any(is_random(i) for i in x) + +def rv(symbol, cls, args, **kwargs): + args = list(map(sympify, args)) + dist = cls(*args) + if kwargs.pop('check', True): + dist.check(*args) + pspace = SingleContinuousPSpace(symbol, dist) + if any(is_random(arg) for arg in args): + from sympy.stats.compound_rv import CompoundPSpace, CompoundDistribution + pspace = CompoundPSpace(symbol, CompoundDistribution(dist)) + return pspace.value + + +class ContinuousDistributionHandmade(SingleContinuousDistribution): + _argnames = ('pdf',) + + def __new__(cls, pdf, set=Interval(-oo, oo)): + return Basic.__new__(cls, pdf, set) + + @property + def set(self): + return self.args[1] + + @staticmethod + def check(pdf, set): + x = Dummy('x') + val = integrate(pdf(x), (x, set)) + _value_check(Eq(val, 1) != S.false, "The pdf on the given set is incorrect.") + + +def ContinuousRV(symbol, density, set=Interval(-oo, oo), **kwargs): + """ + Create a Continuous Random Variable given the following: + + Parameters + ========== + + symbol : Symbol + Represents name of the random variable. + density : Expression containing symbol + Represents probability density function. + set : set/Interval + Represents the region where the pdf is valid, by default is real line. + check : bool + If True, it will check whether the given density + integrates to 1 over the given set. If False, it + will not perform this check. Default is False. + + + Returns + ======= + + RandomSymbol + + Many common continuous random variable types are already implemented. + This function should be necessary only very rarely. + + + Examples + ======== + + >>> from sympy import Symbol, sqrt, exp, pi + >>> from sympy.stats import ContinuousRV, P, E + + >>> x = Symbol("x") + + >>> pdf = sqrt(2)*exp(-x**2/2)/(2*sqrt(pi)) # Normal distribution + >>> X = ContinuousRV(x, pdf) + + >>> E(X) + 0 + >>> P(X>0) + 1/2 + """ + pdf = Piecewise((density, set.as_relational(symbol)), (0, True)) + pdf = Lambda(symbol, pdf) + # have a default of False while `rv` should have a default of True + kwargs['check'] = kwargs.pop('check', False) + return rv(symbol.name, ContinuousDistributionHandmade, (pdf, set), **kwargs) + +######################################## +# Continuous Probability Distributions # +######################################## + +#------------------------------------------------------------------------------- +# Arcsin distribution ---------------------------------------------------------- + + +class ArcsinDistribution(SingleContinuousDistribution): + _argnames = ('a', 'b') + + @property + def set(self): + return Interval(self.a, self.b) + + def pdf(self, x): + a, b = self.a, self.b + return 1/(pi*sqrt((x - a)*(b - x))) + + def _cdf(self, x): + a, b = self.a, self.b + return Piecewise( + (S.Zero, x < a), + (2*asin(sqrt((x - a)/(b - a)))/pi, x <= b), + (S.One, True)) + + +def Arcsin(name, a=0, b=1): + r""" + Create a Continuous Random Variable with an arcsin distribution. + + The density of the arcsin distribution is given by + + .. math:: + f(x) := \frac{1}{\pi\sqrt{(x-a)(b-x)}} + + with :math:`x \in (a,b)`. It must hold that :math:`-\infty < a < b < \infty`. + + Parameters + ========== + + a : Real number, the left interval boundary + b : Real number, the right interval boundary + + Returns + ======= + + RandomSymbol + + Examples + ======== + + >>> from sympy.stats import Arcsin, density, cdf + >>> from sympy import Symbol + + >>> a = Symbol("a", real=True) + >>> b = Symbol("b", real=True) + >>> z = Symbol("z") + + >>> X = Arcsin("x", a, b) + + >>> density(X)(z) + 1/(pi*sqrt((-a + z)*(b - z))) + + >>> cdf(X)(z) + Piecewise((0, a > z), + (2*asin(sqrt((-a + z)/(-a + b)))/pi, b >= z), + (1, True)) + + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Arcsine_distribution + + """ + + return rv(name, ArcsinDistribution, (a, b)) + +#------------------------------------------------------------------------------- +# Benini distribution ---------------------------------------------------------- + + +class BeniniDistribution(SingleContinuousDistribution): + _argnames = ('alpha', 'beta', 'sigma') + + @staticmethod + def check(alpha, beta, sigma): + _value_check(alpha > 0, "Shape parameter Alpha must be positive.") + _value_check(beta > 0, "Shape parameter Beta must be positive.") + _value_check(sigma > 0, "Scale parameter Sigma must be positive.") + + @property + def set(self): + return Interval(self.sigma, oo) + + def pdf(self, x): + alpha, beta, sigma = self.alpha, self.beta, self.sigma + return (exp(-alpha*log(x/sigma) - beta*log(x/sigma)**2) + *(alpha/x + 2*beta*log(x/sigma)/x)) + + def _moment_generating_function(self, t): + raise NotImplementedError('The moment generating function of the ' + 'Benini distribution does not exist.') + +def Benini(name, alpha, beta, sigma): + r""" + Create a Continuous Random Variable with a Benini distribution. + + The density of the Benini distribution is given by + + .. math:: + f(x) := e^{-\alpha\log{\frac{x}{\sigma}} + -\beta\log^2\left[{\frac{x}{\sigma}}\right]} + \left(\frac{\alpha}{x}+\frac{2\beta\log{\frac{x}{\sigma}}}{x}\right) + + This is a heavy-tailed distribution and is also known as the log-Rayleigh + distribution. + + Parameters + ========== + + alpha : Real number, `\alpha > 0`, a shape + beta : Real number, `\beta > 0`, a shape + sigma : Real number, `\sigma > 0`, a scale + + Returns + ======= + + RandomSymbol + + Examples + ======== + + >>> from sympy.stats import Benini, density, cdf + >>> from sympy import Symbol, pprint + + >>> alpha = Symbol("alpha", positive=True) + >>> beta = Symbol("beta", positive=True) + >>> sigma = Symbol("sigma", positive=True) + >>> z = Symbol("z") + + >>> X = Benini("x", alpha, beta, sigma) + + >>> D = density(X)(z) + >>> pprint(D, use_unicode=False) + / / z \\ / z \ 2/ z \ + | 2*beta*log|-----|| - alpha*log|-----| - beta*log |-----| + |alpha \sigma/| \sigma/ \sigma/ + |----- + -----------------|*e + \ z z / + + >>> cdf(X)(z) + Piecewise((1 - exp(-alpha*log(z/sigma) - beta*log(z/sigma)**2), sigma <= z), + (0, True)) + + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Benini_distribution + .. [2] https://reference.wolfram.com/legacy/v8/ref/BeniniDistribution.html + + """ + + return rv(name, BeniniDistribution, (alpha, beta, sigma)) + +#------------------------------------------------------------------------------- +# Beta distribution ------------------------------------------------------------ + + +class BetaDistribution(SingleContinuousDistribution): + _argnames = ('alpha', 'beta') + + set = Interval(0, 1) + + @staticmethod + def check(alpha, beta): + _value_check(alpha > 0, "Shape parameter Alpha must be positive.") + _value_check(beta > 0, "Shape parameter Beta must be positive.") + + def pdf(self, x): + alpha, beta = self.alpha, self.beta + return x**(alpha - 1) * (1 - x)**(beta - 1) / beta_fn(alpha, beta) + + def _characteristic_function(self, t): + return hyper((self.alpha,), (self.alpha + self.beta,), I*t) + + def _moment_generating_function(self, t): + return hyper((self.alpha,), (self.alpha + self.beta,), t) + + +def Beta(name, alpha, beta): + r""" + Create a Continuous Random Variable with a Beta distribution. + + The density of the Beta distribution is given by + + .. math:: + f(x) := \frac{x^{\alpha-1}(1-x)^{\beta-1}} {\mathrm{B}(\alpha,\beta)} + + with :math:`x \in [0,1]`. + + Parameters + ========== + + alpha : Real number, `\alpha > 0`, a shape + beta : Real number, `\beta > 0`, a shape + + Returns + ======= + + RandomSymbol + + Examples + ======== + + >>> from sympy.stats import Beta, density, E, variance + >>> from sympy import Symbol, simplify, pprint, factor + + >>> alpha = Symbol("alpha", positive=True) + >>> beta = Symbol("beta", positive=True) + >>> z = Symbol("z") + + >>> X = Beta("x", alpha, beta) + + >>> D = density(X)(z) + >>> pprint(D, use_unicode=False) + alpha - 1 beta - 1 + z *(1 - z) + -------------------------- + B(alpha, beta) + + >>> simplify(E(X)) + alpha/(alpha + beta) + + >>> factor(simplify(variance(X))) + alpha*beta/((alpha + beta)**2*(alpha + beta + 1)) + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Beta_distribution + .. [2] https://mathworld.wolfram.com/BetaDistribution.html + + """ + + return rv(name, BetaDistribution, (alpha, beta)) + +#------------------------------------------------------------------------------- +# Noncentral Beta distribution ------------------------------------------------------------ + + +class BetaNoncentralDistribution(SingleContinuousDistribution): + _argnames = ('alpha', 'beta', 'lamda') + + set = Interval(0, 1) + + @staticmethod + def check(alpha, beta, lamda): + _value_check(alpha > 0, "Shape parameter Alpha must be positive.") + _value_check(beta > 0, "Shape parameter Beta must be positive.") + _value_check(lamda >= 0, "Noncentrality parameter Lambda must be positive") + + def pdf(self, x): + alpha, beta, lamda = self.alpha, self.beta, self.lamda + k = Dummy("k") + return Sum(exp(-lamda / 2) * (lamda / 2)**k * x**(alpha + k - 1) *( + 1 - x)**(beta - 1) / (factorial(k) * beta_fn(alpha + k, beta)), (k, 0, oo)) + +def BetaNoncentral(name, alpha, beta, lamda): + r""" + Create a Continuous Random Variable with a Type I Noncentral Beta distribution. + + The density of the Noncentral Beta distribution is given by + + .. math:: + f(x) := \sum_{k=0}^\infty e^{-\lambda/2}\frac{(\lambda/2)^k}{k!} + \frac{x^{\alpha+k-1}(1-x)^{\beta-1}}{\mathrm{B}(\alpha+k,\beta)} + + with :math:`x \in [0,1]`. + + Parameters + ========== + + alpha : Real number, `\alpha > 0`, a shape + beta : Real number, `\beta > 0`, a shape + lamda : Real number, `\lambda \geq 0`, noncentrality parameter + + Returns + ======= + + RandomSymbol + + Examples + ======== + + >>> from sympy.stats import BetaNoncentral, density, cdf + >>> from sympy import Symbol, pprint + + >>> alpha = Symbol("alpha", positive=True) + >>> beta = Symbol("beta", positive=True) + >>> lamda = Symbol("lamda", nonnegative=True) + >>> z = Symbol("z") + + >>> X = BetaNoncentral("x", alpha, beta, lamda) + + >>> D = density(X)(z) + >>> pprint(D, use_unicode=False) + oo + _____ + \ ` + \ -lamda + \ k ------- + \ k + alpha - 1 /lamda\ beta - 1 2 + ) z *|-----| *(1 - z) *e + / \ 2 / + / ------------------------------------------------ + / B(k + alpha, beta)*k! + /____, + k = 0 + + Compute cdf with specific 'x', 'alpha', 'beta' and 'lamda' values as follows: + + >>> cdf(BetaNoncentral("x", 1, 1, 1), evaluate=False)(2).doit() + 2*exp(1/2) + + The argument evaluate=False prevents an attempt at evaluation + of the sum for general x, before the argument 2 is passed. + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Noncentral_beta_distribution + .. [2] https://reference.wolfram.com/language/ref/NoncentralBetaDistribution.html + + """ + + return rv(name, BetaNoncentralDistribution, (alpha, beta, lamda)) + + +#------------------------------------------------------------------------------- +# Beta prime distribution ------------------------------------------------------ + + +class BetaPrimeDistribution(SingleContinuousDistribution): + _argnames = ('alpha', 'beta') + + @staticmethod + def check(alpha, beta): + _value_check(alpha > 0, "Shape parameter Alpha must be positive.") + _value_check(beta > 0, "Shape parameter Beta must be positive.") + + set = Interval(0, oo) + + def pdf(self, x): + alpha, beta = self.alpha, self.beta + return x**(alpha - 1)*(1 + x)**(-alpha - beta)/beta_fn(alpha, beta) + +def BetaPrime(name, alpha, beta): + r""" + Create a continuous random variable with a Beta prime distribution. + + The density of the Beta prime distribution is given by + + .. math:: + f(x) := \frac{x^{\alpha-1} (1+x)^{-\alpha -\beta}}{B(\alpha,\beta)} + + with :math:`x > 0`. + + Parameters + ========== + + alpha : Real number, `\alpha > 0`, a shape + beta : Real number, `\beta > 0`, a shape + + Returns + ======= + + RandomSymbol + + Examples + ======== + + >>> from sympy.stats import BetaPrime, density + >>> from sympy import Symbol, pprint + + >>> alpha = Symbol("alpha", positive=True) + >>> beta = Symbol("beta", positive=True) + >>> z = Symbol("z") + + >>> X = BetaPrime("x", alpha, beta) + + >>> D = density(X)(z) + >>> pprint(D, use_unicode=False) + alpha - 1 -alpha - beta + z *(z + 1) + ------------------------------- + B(alpha, beta) + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Beta_prime_distribution + .. [2] https://mathworld.wolfram.com/BetaPrimeDistribution.html + + """ + + return rv(name, BetaPrimeDistribution, (alpha, beta)) + +#------------------------------------------------------------------------------- +# Bounded Pareto Distribution -------------------------------------------------- +class BoundedParetoDistribution(SingleContinuousDistribution): + _argnames = ('alpha', 'left', 'right') + + @property + def set(self): + return Interval(self.left, self.right) + + @staticmethod + def check(alpha, left, right): + _value_check (alpha.is_positive, "Shape must be positive.") + _value_check (left.is_positive, "Left value should be positive.") + _value_check (right > left, "Right should be greater than left.") + + def pdf(self, x): + alpha, left, right = self.alpha, self.left, self.right + num = alpha * (left**alpha) * x**(- alpha -1) + den = 1 - (left/right)**alpha + return num/den + +def BoundedPareto(name, alpha, left, right): + r""" + Create a continuous random variable with a Bounded Pareto distribution. + + The density of the Bounded Pareto distribution is given by + + .. math:: + f(x) := \frac{\alpha L^{\alpha}x^{-\alpha-1}}{1-(\frac{L}{H})^{\alpha}} + + Parameters + ========== + + alpha : Real Number, `\alpha > 0` + Shape parameter + left : Real Number, `left > 0` + Location parameter + right : Real Number, `right > left` + Location parameter + + Examples + ======== + + >>> from sympy.stats import BoundedPareto, density, cdf, E + >>> from sympy import symbols + >>> L, H = symbols('L, H', positive=True) + >>> X = BoundedPareto('X', 2, L, H) + >>> x = symbols('x') + >>> density(X)(x) + 2*L**2/(x**3*(1 - L**2/H**2)) + >>> cdf(X)(x) + Piecewise((-H**2*L**2/(x**2*(H**2 - L**2)) + H**2/(H**2 - L**2), L <= x), (0, True)) + >>> E(X).simplify() + 2*H*L/(H + L) + + Returns + ======= + + RandomSymbol + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Pareto_distribution#Bounded_Pareto_distribution + + """ + return rv (name, BoundedParetoDistribution, (alpha, left, right)) + +# ------------------------------------------------------------------------------ +# Cauchy distribution ---------------------------------------------------------- + + +class CauchyDistribution(SingleContinuousDistribution): + _argnames = ('x0', 'gamma') + + @staticmethod + def check(x0, gamma): + _value_check(gamma > 0, "Scale parameter Gamma must be positive.") + _value_check(x0.is_real, "Location parameter must be real.") + + def pdf(self, x): + return 1/(pi*self.gamma*(1 + ((x - self.x0)/self.gamma)**2)) + + def _cdf(self, x): + x0, gamma = self.x0, self.gamma + return (1/pi)*atan((x - x0)/gamma) + S.Half + + def _characteristic_function(self, t): + return exp(self.x0 * I * t - self.gamma * Abs(t)) + + def _moment_generating_function(self, t): + raise NotImplementedError("The moment generating function for the " + "Cauchy distribution does not exist.") + + def _quantile(self, p): + return self.x0 + self.gamma*tan(pi*(p - S.Half)) + + +def Cauchy(name, x0, gamma): + r""" + Create a continuous random variable with a Cauchy distribution. + + The density of the Cauchy distribution is given by + + .. math:: + f(x) := \frac{1}{\pi \gamma [1 + {(\frac{x-x_0}{\gamma})}^2]} + + Parameters + ========== + + x0 : Real number, the location + gamma : Real number, `\gamma > 0`, a scale + + Returns + ======= + + RandomSymbol + + Examples + ======== + + >>> from sympy.stats import Cauchy, density + >>> from sympy import Symbol + + >>> x0 = Symbol("x0") + >>> gamma = Symbol("gamma", positive=True) + >>> z = Symbol("z") + + >>> X = Cauchy("x", x0, gamma) + + >>> density(X)(z) + 1/(pi*gamma*(1 + (-x0 + z)**2/gamma**2)) + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Cauchy_distribution + .. [2] https://mathworld.wolfram.com/CauchyDistribution.html + + """ + + return rv(name, CauchyDistribution, (x0, gamma)) + +#------------------------------------------------------------------------------- +# Chi distribution ------------------------------------------------------------- + + +class ChiDistribution(SingleContinuousDistribution): + _argnames = ('k',) + + @staticmethod + def check(k): + _value_check(k > 0, "Number of degrees of freedom (k) must be positive.") + _value_check(k.is_integer, "Number of degrees of freedom (k) must be an integer.") + + set = Interval(0, oo) + + def pdf(self, x): + return 2**(1 - self.k/2)*x**(self.k - 1)*exp(-x**2/2)/gamma(self.k/2) + + def _characteristic_function(self, t): + k = self.k + + part_1 = hyper((k/2,), (S.Half,), -t**2/2) + part_2 = I*t*sqrt(2)*gamma((k+1)/2)/gamma(k/2) + part_3 = hyper(((k+1)/2,), (Rational(3, 2),), -t**2/2) + return part_1 + part_2*part_3 + + def _moment_generating_function(self, t): + k = self.k + + part_1 = hyper((k / 2,), (S.Half,), t ** 2 / 2) + part_2 = t * sqrt(2) * gamma((k + 1) / 2) / gamma(k / 2) + part_3 = hyper(((k + 1) / 2,), (S(3) / 2,), t ** 2 / 2) + return part_1 + part_2 * part_3 + +def Chi(name, k): + r""" + Create a continuous random variable with a Chi distribution. + + The density of the Chi distribution is given by + + .. math:: + f(x) := \frac{2^{1-k/2}x^{k-1}e^{-x^2/2}}{\Gamma(k/2)} + + with :math:`x \geq 0`. + + Parameters + ========== + + k : Positive integer, The number of degrees of freedom + + Returns + ======= + + RandomSymbol + + Examples + ======== + + >>> from sympy.stats import Chi, density, E + >>> from sympy import Symbol, simplify + + >>> k = Symbol("k", integer=True) + >>> z = Symbol("z") + + >>> X = Chi("x", k) + + >>> density(X)(z) + 2**(1 - k/2)*z**(k - 1)*exp(-z**2/2)/gamma(k/2) + + >>> simplify(E(X)) + sqrt(2)*gamma(k/2 + 1/2)/gamma(k/2) + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Chi_distribution + .. [2] https://mathworld.wolfram.com/ChiDistribution.html + + """ + + return rv(name, ChiDistribution, (k,)) + +#------------------------------------------------------------------------------- +# Non-central Chi distribution ------------------------------------------------- + + +class ChiNoncentralDistribution(SingleContinuousDistribution): + _argnames = ('k', 'l') + + @staticmethod + def check(k, l): + _value_check(k > 0, "Number of degrees of freedom (k) must be positive.") + _value_check(k.is_integer, "Number of degrees of freedom (k) must be an integer.") + _value_check(l > 0, "Shift parameter Lambda must be positive.") + + set = Interval(0, oo) + + def pdf(self, x): + k, l = self.k, self.l + return exp(-(x**2+l**2)/2)*x**k*l / (l*x)**(k/2) * besseli(k/2-1, l*x) + +def ChiNoncentral(name, k, l): + r""" + Create a continuous random variable with a non-central Chi distribution. + + Explanation + =========== + + The density of the non-central Chi distribution is given by + + .. math:: + f(x) := \frac{e^{-(x^2+\lambda^2)/2} x^k\lambda} + {(\lambda x)^{k/2}} I_{k/2-1}(\lambda x) + + with `x \geq 0`. Here, `I_\nu (x)` is the + :ref:`modified Bessel function of the first kind `. + + Parameters + ========== + + k : A positive Integer, $k > 0$ + The number of degrees of freedom. + lambda : Real number, `\lambda > 0` + Shift parameter. + + Returns + ======= + + RandomSymbol + + Examples + ======== + + >>> from sympy.stats import ChiNoncentral, density + >>> from sympy import Symbol + + >>> k = Symbol("k", integer=True) + >>> l = Symbol("l") + >>> z = Symbol("z") + + >>> X = ChiNoncentral("x", k, l) + + >>> density(X)(z) + l*z**k*exp(-l**2/2 - z**2/2)*besseli(k/2 - 1, l*z)/(l*z)**(k/2) + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Noncentral_chi_distribution + """ + + return rv(name, ChiNoncentralDistribution, (k, l)) + +#------------------------------------------------------------------------------- +# Chi squared distribution ----------------------------------------------------- + + +class ChiSquaredDistribution(SingleContinuousDistribution): + _argnames = ('k',) + + @staticmethod + def check(k): + _value_check(k > 0, "Number of degrees of freedom (k) must be positive.") + _value_check(k.is_integer, "Number of degrees of freedom (k) must be an integer.") + + set = Interval(0, oo) + + def pdf(self, x): + k = self.k + return 1/(2**(k/2)*gamma(k/2))*x**(k/2 - 1)*exp(-x/2) + + def _cdf(self, x): + k = self.k + return Piecewise( + (S.One/gamma(k/2)*lowergamma(k/2, x/2), x >= 0), + (0, True) + ) + + def _characteristic_function(self, t): + return (1 - 2*I*t)**(-self.k/2) + + def _moment_generating_function(self, t): + return (1 - 2*t)**(-self.k/2) + + +def ChiSquared(name, k): + r""" + Create a continuous random variable with a Chi-squared distribution. + + Explanation + =========== + + The density of the Chi-squared distribution is given by + + .. math:: + f(x) := \frac{1}{2^{\frac{k}{2}}\Gamma\left(\frac{k}{2}\right)} + x^{\frac{k}{2}-1} e^{-\frac{x}{2}} + + with :math:`x \geq 0`. + + Parameters + ========== + + k : Positive integer + The number of degrees of freedom. + + Returns + ======= + + RandomSymbol + + Examples + ======== + + >>> from sympy.stats import ChiSquared, density, E, variance, moment + >>> from sympy import Symbol + + >>> k = Symbol("k", integer=True, positive=True) + >>> z = Symbol("z") + + >>> X = ChiSquared("x", k) + + >>> density(X)(z) + z**(k/2 - 1)*exp(-z/2)/(2**(k/2)*gamma(k/2)) + + >>> E(X) + k + + >>> variance(X) + 2*k + + >>> moment(X, 3) + k**3 + 6*k**2 + 8*k + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Chi_squared_distribution + .. [2] https://mathworld.wolfram.com/Chi-SquaredDistribution.html + """ + + return rv(name, ChiSquaredDistribution, (k, )) + +#------------------------------------------------------------------------------- +# Dagum distribution ----------------------------------------------------------- + + +class DagumDistribution(SingleContinuousDistribution): + _argnames = ('p', 'a', 'b') + + set = Interval(0, oo) + + @staticmethod + def check(p, a, b): + _value_check(p > 0, "Shape parameter p must be positive.") + _value_check(a > 0, "Shape parameter a must be positive.") + _value_check(b > 0, "Scale parameter b must be positive.") + + def pdf(self, x): + p, a, b = self.p, self.a, self.b + return a*p/x*((x/b)**(a*p)/(((x/b)**a + 1)**(p + 1))) + + def _cdf(self, x): + p, a, b = self.p, self.a, self.b + return Piecewise(((S.One + (S(x)/b)**-a)**-p, x>=0), + (S.Zero, True)) + +def Dagum(name, p, a, b): + r""" + Create a continuous random variable with a Dagum distribution. + + Explanation + =========== + + The density of the Dagum distribution is given by + + .. math:: + f(x) := \frac{a p}{x} \left( \frac{\left(\tfrac{x}{b}\right)^{a p}} + {\left(\left(\tfrac{x}{b}\right)^a + 1 \right)^{p+1}} \right) + + with :math:`x > 0`. + + Parameters + ========== + + p : Real number + `p > 0`, a shape. + a : Real number + `a > 0`, a shape. + b : Real number + `b > 0`, a scale. + + Returns + ======= + + RandomSymbol + + Examples + ======== + + >>> from sympy.stats import Dagum, density, cdf + >>> from sympy import Symbol + + >>> p = Symbol("p", positive=True) + >>> a = Symbol("a", positive=True) + >>> b = Symbol("b", positive=True) + >>> z = Symbol("z") + + >>> X = Dagum("x", p, a, b) + + >>> density(X)(z) + a*p*(z/b)**(a*p)*((z/b)**a + 1)**(-p - 1)/z + + >>> cdf(X)(z) + Piecewise(((1 + (z/b)**(-a))**(-p), z >= 0), (0, True)) + + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Dagum_distribution + + """ + + return rv(name, DagumDistribution, (p, a, b)) + +#------------------------------------------------------------------------------- +# Davis distribution ----------------------------------------------------------- + +class DavisDistribution(SingleContinuousDistribution): + _argnames = ('b', 'n', 'mu') + + set = Interval(0, oo) + + @staticmethod + def check(b, n, mu): + _value_check(b > 0, "Scale parameter b must be positive.") + _value_check(n > 1, "Shape parameter n must be above 1.") + _value_check(mu > 0, "Location parameter mu must be positive.") + + def pdf(self, x): + b, n, mu = self.b, self.n, self.mu + dividend = b**n*(x - mu)**(-1-n) + divisor = (exp(b/(x-mu))-1)*(gamma(n)*zeta(n)) + return dividend/divisor + + +def Davis(name, b, n, mu): + r""" Create a continuous random variable with Davis distribution. + + Explanation + =========== + + The density of Davis distribution is given by + + .. math:: + f(x; \mu; b, n) := \frac{b^{n}(x - \mu)^{1-n}}{ \left( e^{\frac{b}{x-\mu}} - 1 \right) \Gamma(n)\zeta(n)} + + with :math:`x \in [0,\infty]`. + + Davis distribution is a generalization of the Planck's law of radiation from statistical physics. It is used for modeling income distribution. + + Parameters + ========== + b : Real number + `p > 0`, a scale. + n : Real number + `n > 1`, a shape. + mu : Real number + `mu > 0`, a location. + + Returns + ======= + + RandomSymbol + + Examples + ======== + >>> from sympy.stats import Davis, density + >>> from sympy import Symbol + >>> b = Symbol("b", positive=True) + >>> n = Symbol("n", positive=True) + >>> mu = Symbol("mu", positive=True) + >>> z = Symbol("z") + >>> X = Davis("x", b, n, mu) + >>> density(X)(z) + b**n*(-mu + z)**(-n - 1)/((exp(b/(-mu + z)) - 1)*gamma(n)*zeta(n)) + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Davis_distribution + .. [2] https://reference.wolfram.com/language/ref/DavisDistribution.html + + """ + return rv(name, DavisDistribution, (b, n, mu)) + + +#------------------------------------------------------------------------------- +# Erlang distribution ---------------------------------------------------------- + + +def Erlang(name, k, l): + r""" + Create a continuous random variable with an Erlang distribution. + + Explanation + =========== + + The density of the Erlang distribution is given by + + .. math:: + f(x) := \frac{\lambda^k x^{k-1} e^{-\lambda x}}{(k-1)!} + + with :math:`x \in [0,\infty]`. + + Parameters + ========== + + k : Positive integer + l : Real number, `\lambda > 0`, the rate + + Returns + ======= + + RandomSymbol + + Examples + ======== + + >>> from sympy.stats import Erlang, density, cdf, E, variance + >>> from sympy import Symbol, simplify, pprint + + >>> k = Symbol("k", integer=True, positive=True) + >>> l = Symbol("l", positive=True) + >>> z = Symbol("z") + + >>> X = Erlang("x", k, l) + + >>> D = density(X)(z) + >>> pprint(D, use_unicode=False) + k k - 1 -l*z + l *z *e + --------------- + Gamma(k) + + >>> C = cdf(X)(z) + >>> pprint(C, use_unicode=False) + /lowergamma(k, l*z) + |------------------ for z > 0 + < Gamma(k) + | + \ 0 otherwise + + + >>> E(X) + k/l + + >>> simplify(variance(X)) + k/l**2 + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Erlang_distribution + .. [2] https://mathworld.wolfram.com/ErlangDistribution.html + + """ + + return rv(name, GammaDistribution, (k, S.One/l)) + +# ------------------------------------------------------------------------------- +# ExGaussian distribution ----------------------------------------------------- + + +class ExGaussianDistribution(SingleContinuousDistribution): + _argnames = ('mean', 'std', 'rate') + + set = Interval(-oo, oo) + + @staticmethod + def check(mean, std, rate): + _value_check( + std > 0, "Standard deviation of ExGaussian must be positive.") + _value_check(rate > 0, "Rate of ExGaussian must be positive.") + + def pdf(self, x): + mean, std, rate = self.mean, self.std, self.rate + term1 = rate/2 + term2 = exp(rate * (2 * mean + rate * std**2 - 2*x)/2) + term3 = erfc((mean + rate*std**2 - x)/(sqrt(2)*std)) + return term1*term2*term3 + + def _cdf(self, x): + from sympy.stats import cdf + mean, std, rate = self.mean, self.std, self.rate + u = rate*(x - mean) + v = rate*std + GaussianCDF1 = cdf(Normal('x', 0, v))(u) + GaussianCDF2 = cdf(Normal('x', v**2, v))(u) + + return GaussianCDF1 - exp(-u + (v**2/2) + log(GaussianCDF2)) + + def _characteristic_function(self, t): + mean, std, rate = self.mean, self.std, self.rate + term1 = (1 - I*t/rate)**(-1) + term2 = exp(I*mean*t - std**2*t**2/2) + return term1 * term2 + + def _moment_generating_function(self, t): + mean, std, rate = self.mean, self.std, self.rate + term1 = (1 - t/rate)**(-1) + term2 = exp(mean*t + std**2*t**2/2) + return term1*term2 + + +def ExGaussian(name, mean, std, rate): + r""" + Create a continuous random variable with an Exponentially modified + Gaussian (EMG) distribution. + + Explanation + =========== + + The density of the exponentially modified Gaussian distribution is given by + + .. math:: + f(x) := \frac{\lambda}{2}e^{\frac{\lambda}{2}(2\mu+\lambda\sigma^2-2x)} + \text{erfc}(\frac{\mu + \lambda\sigma^2 - x}{\sqrt{2}\sigma}) + + with $x > 0$. Note that the expected value is `1/\lambda`. + + Parameters + ========== + + name : A string giving a name for this distribution + mean : A Real number, the mean of Gaussian component + std : A positive Real number, + :math: `\sigma^2 > 0` the variance of Gaussian component + rate : A positive Real number, + :math: `\lambda > 0` the rate of Exponential component + + Returns + ======= + + RandomSymbol + + Examples + ======== + + >>> from sympy.stats import ExGaussian, density, cdf, E + >>> from sympy.stats import variance, skewness + >>> from sympy import Symbol, pprint, simplify + + >>> mean = Symbol("mu") + >>> std = Symbol("sigma", positive=True) + >>> rate = Symbol("lamda", positive=True) + >>> z = Symbol("z") + >>> X = ExGaussian("x", mean, std, rate) + + >>> pprint(density(X)(z), use_unicode=False) + / 2 \ + lamda*\lamda*sigma + 2*mu - 2*z/ + --------------------------------- / ___ / 2 \\ + 2 |\/ 2 *\lamda*sigma + mu - z/| + lamda*e *erfc|-----------------------------| + \ 2*sigma / + ---------------------------------------------------------------------------- + 2 + + >>> cdf(X)(z) + -(erf(sqrt(2)*(-lamda**2*sigma**2 + lamda*(-mu + z))/(2*lamda*sigma))/2 + 1/2)*exp(lamda**2*sigma**2/2 - lamda*(-mu + z)) + erf(sqrt(2)*(-mu + z)/(2*sigma))/2 + 1/2 + + >>> E(X) + (lamda*mu + 1)/lamda + + >>> simplify(variance(X)) + sigma**2 + lamda**(-2) + + >>> simplify(skewness(X)) + 2/(lamda**2*sigma**2 + 1)**(3/2) + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Exponentially_modified_Gaussian_distribution + """ + return rv(name, ExGaussianDistribution, (mean, std, rate)) + +#------------------------------------------------------------------------------- +# Exponential distribution ----------------------------------------------------- + + +class ExponentialDistribution(SingleContinuousDistribution): + _argnames = ('rate',) + + set = Interval(0, oo) + + @staticmethod + def check(rate): + _value_check(rate > 0, "Rate must be positive.") + + def pdf(self, x): + return self.rate * exp(-self.rate*x) + + def _cdf(self, x): + return Piecewise( + (S.One - exp(-self.rate*x), x >= 0), + (0, True), + ) + + def _characteristic_function(self, t): + rate = self.rate + return rate / (rate - I*t) + + def _moment_generating_function(self, t): + rate = self.rate + return rate / (rate - t) + + def _quantile(self, p): + return -log(1-p)/self.rate + + +def Exponential(name, rate): + r""" + Create a continuous random variable with an Exponential distribution. + + Explanation + =========== + + The density of the exponential distribution is given by + + .. math:: + f(x) := \lambda \exp(-\lambda x) + + with $x > 0$. Note that the expected value is `1/\lambda`. + + Parameters + ========== + + rate : A positive Real number, `\lambda > 0`, the rate (or inverse scale/inverse mean) + + Returns + ======= + + RandomSymbol + + Examples + ======== + + >>> from sympy.stats import Exponential, density, cdf, E + >>> from sympy.stats import variance, std, skewness, quantile + >>> from sympy import Symbol + + >>> l = Symbol("lambda", positive=True) + >>> z = Symbol("z") + >>> p = Symbol("p") + >>> X = Exponential("x", l) + + >>> density(X)(z) + lambda*exp(-lambda*z) + + >>> cdf(X)(z) + Piecewise((1 - exp(-lambda*z), z >= 0), (0, True)) + + >>> quantile(X)(p) + -log(1 - p)/lambda + + >>> E(X) + 1/lambda + + >>> variance(X) + lambda**(-2) + + >>> skewness(X) + 2 + + >>> X = Exponential('x', 10) + + >>> density(X)(z) + 10*exp(-10*z) + + >>> E(X) + 1/10 + + >>> std(X) + 1/10 + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Exponential_distribution + .. [2] https://mathworld.wolfram.com/ExponentialDistribution.html + + """ + + return rv(name, ExponentialDistribution, (rate, )) + + +# ------------------------------------------------------------------------------- +# Exponential Power distribution ----------------------------------------------------- + +class ExponentialPowerDistribution(SingleContinuousDistribution): + _argnames = ('mu', 'alpha', 'beta') + + set = Interval(-oo, oo) + + @staticmethod + def check(mu, alpha, beta): + _value_check(alpha > 0, "Scale parameter alpha must be positive.") + _value_check(beta > 0, "Shape parameter beta must be positive.") + + def pdf(self, x): + mu, alpha, beta = self.mu, self.alpha, self.beta + num = beta*exp(-(Abs(x - mu)/alpha)**beta) + den = 2*alpha*gamma(1/beta) + return num/den + + def _cdf(self, x): + mu, alpha, beta = self.mu, self.alpha, self.beta + num = lowergamma(1/beta, (Abs(x - mu) / alpha)**beta) + den = 2*gamma(1/beta) + return sign(x - mu)*num/den + S.Half + + +def ExponentialPower(name, mu, alpha, beta): + r""" + Create a Continuous Random Variable with Exponential Power distribution. + This distribution is known also as Generalized Normal + distribution version 1. + + Explanation + =========== + + The density of the Exponential Power distribution is given by + + .. math:: + f(x) := \frac{\beta}{2\alpha\Gamma(\frac{1}{\beta})} + e^{{-(\frac{|x - \mu|}{\alpha})^{\beta}}} + + with :math:`x \in [ - \infty, \infty ]`. + + Parameters + ========== + + mu : Real number + A location. + alpha : Real number,`\alpha > 0` + A scale. + beta : Real number, `\beta > 0` + A shape. + + Returns + ======= + + RandomSymbol + + Examples + ======== + + >>> from sympy.stats import ExponentialPower, density, cdf + >>> from sympy import Symbol, pprint + >>> z = Symbol("z") + >>> mu = Symbol("mu") + >>> alpha = Symbol("alpha", positive=True) + >>> beta = Symbol("beta", positive=True) + >>> X = ExponentialPower("x", mu, alpha, beta) + >>> pprint(density(X)(z), use_unicode=False) + beta + /|mu - z|\ + -|--------| + \ alpha / + beta*e + --------------------- + / 1 \ + 2*alpha*Gamma|----| + \beta/ + >>> cdf(X)(z) + 1/2 + lowergamma(1/beta, (Abs(mu - z)/alpha)**beta)*sign(-mu + z)/(2*gamma(1/beta)) + + References + ========== + + .. [1] https://reference.wolfram.com/language/ref/ExponentialPowerDistribution.html + .. [2] https://en.wikipedia.org/wiki/Generalized_normal_distribution#Version_1 + + """ + return rv(name, ExponentialPowerDistribution, (mu, alpha, beta)) + + +#------------------------------------------------------------------------------- +# F distribution --------------------------------------------------------------- + + +class FDistributionDistribution(SingleContinuousDistribution): + _argnames = ('d1', 'd2') + + set = Interval(0, oo) + + @staticmethod + def check(d1, d2): + _value_check((d1 > 0, d1.is_integer), + "Degrees of freedom d1 must be positive integer.") + _value_check((d2 > 0, d2.is_integer), + "Degrees of freedom d2 must be positive integer.") + + def pdf(self, x): + d1, d2 = self.d1, self.d2 + return (sqrt((d1*x)**d1*d2**d2 / (d1*x+d2)**(d1+d2)) + / (x * beta_fn(d1/2, d2/2))) + + def _moment_generating_function(self, t): + raise NotImplementedError('The moment generating function for the ' + 'F-distribution does not exist.') + +def FDistribution(name, d1, d2): + r""" + Create a continuous random variable with a F distribution. + + Explanation + =========== + + The density of the F distribution is given by + + .. math:: + f(x) := \frac{\sqrt{\frac{(d_1 x)^{d_1} d_2^{d_2}} + {(d_1 x + d_2)^{d_1 + d_2}}}} + {x \mathrm{B} \left(\frac{d_1}{2}, \frac{d_2}{2}\right)} + + with :math:`x > 0`. + + Parameters + ========== + + d1 : `d_1 > 0`, where `d_1` is the degrees of freedom (`n_1 - 1`) + d2 : `d_2 > 0`, where `d_2` is the degrees of freedom (`n_2 - 1`) + + Returns + ======= + + RandomSymbol + + Examples + ======== + + >>> from sympy.stats import FDistribution, density + >>> from sympy import Symbol, pprint + + >>> d1 = Symbol("d1", positive=True) + >>> d2 = Symbol("d2", positive=True) + >>> z = Symbol("z") + + >>> X = FDistribution("x", d1, d2) + + >>> D = density(X)(z) + >>> pprint(D, use_unicode=False) + d2 + -- ______________________________ + 2 / d1 -d1 - d2 + d2 *\/ (d1*z) *(d1*z + d2) + -------------------------------------- + /d1 d2\ + z*B|--, --| + \2 2 / + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/F-distribution + .. [2] https://mathworld.wolfram.com/F-Distribution.html + + """ + + return rv(name, FDistributionDistribution, (d1, d2)) + +#------------------------------------------------------------------------------- +# Fisher Z distribution -------------------------------------------------------- + +class FisherZDistribution(SingleContinuousDistribution): + _argnames = ('d1', 'd2') + + set = Interval(-oo, oo) + + @staticmethod + def check(d1, d2): + _value_check(d1 > 0, "Degree of freedom d1 must be positive.") + _value_check(d2 > 0, "Degree of freedom d2 must be positive.") + + def pdf(self, x): + d1, d2 = self.d1, self.d2 + return (2*d1**(d1/2)*d2**(d2/2) / beta_fn(d1/2, d2/2) * + exp(d1*x) / (d1*exp(2*x)+d2)**((d1+d2)/2)) + +def FisherZ(name, d1, d2): + r""" + Create a Continuous Random Variable with an Fisher's Z distribution. + + Explanation + =========== + + The density of the Fisher's Z distribution is given by + + .. math:: + f(x) := \frac{2d_1^{d_1/2} d_2^{d_2/2}} {\mathrm{B}(d_1/2, d_2/2)} + \frac{e^{d_1z}}{\left(d_1e^{2z}+d_2\right)^{\left(d_1+d_2\right)/2}} + + + .. TODO - What is the difference between these degrees of freedom? + + Parameters + ========== + + d1 : `d_1 > 0` + Degree of freedom. + d2 : `d_2 > 0` + Degree of freedom. + + Returns + ======= + + RandomSymbol + + Examples + ======== + + >>> from sympy.stats import FisherZ, density + >>> from sympy import Symbol, pprint + + >>> d1 = Symbol("d1", positive=True) + >>> d2 = Symbol("d2", positive=True) + >>> z = Symbol("z") + + >>> X = FisherZ("x", d1, d2) + + >>> D = density(X)(z) + >>> pprint(D, use_unicode=False) + d1 d2 + d1 d2 - -- - -- + -- -- 2 2 + 2 2 / 2*z \ d1*z + 2*d1 *d2 *\d1*e + d2/ *e + ----------------------------------------- + /d1 d2\ + B|--, --| + \2 2 / + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Fisher%27s_z-distribution + .. [2] https://mathworld.wolfram.com/Fishersz-Distribution.html + + """ + + return rv(name, FisherZDistribution, (d1, d2)) + +#------------------------------------------------------------------------------- +# Frechet distribution --------------------------------------------------------- + +class FrechetDistribution(SingleContinuousDistribution): + _argnames = ('a', 's', 'm') + + set = Interval(0, oo) + + @staticmethod + def check(a, s, m): + _value_check(a > 0, "Shape parameter alpha must be positive.") + _value_check(s > 0, "Scale parameter s must be positive.") + + def __new__(cls, a, s=1, m=0): + a, s, m = list(map(sympify, (a, s, m))) + return Basic.__new__(cls, a, s, m) + + def pdf(self, x): + a, s, m = self.a, self.s, self.m + return a/s * ((x-m)/s)**(-1-a) * exp(-((x-m)/s)**(-a)) + + def _cdf(self, x): + a, s, m = self.a, self.s, self.m + return Piecewise((exp(-((x-m)/s)**(-a)), x >= m), + (S.Zero, True)) + +def Frechet(name, a, s=1, m=0): + r""" + Create a continuous random variable with a Frechet distribution. + + Explanation + =========== + + The density of the Frechet distribution is given by + + .. math:: + f(x) := \frac{\alpha}{s} \left(\frac{x-m}{s}\right)^{-1-\alpha} + e^{-(\frac{x-m}{s})^{-\alpha}} + + with :math:`x \geq m`. + + Parameters + ========== + + a : Real number, :math:`a \in \left(0, \infty\right)` the shape + s : Real number, :math:`s \in \left(0, \infty\right)` the scale + m : Real number, :math:`m \in \left(-\infty, \infty\right)` the minimum + + Returns + ======= + + RandomSymbol + + Examples + ======== + + >>> from sympy.stats import Frechet, density, cdf + >>> from sympy import Symbol + + >>> a = Symbol("a", positive=True) + >>> s = Symbol("s", positive=True) + >>> m = Symbol("m", real=True) + >>> z = Symbol("z") + + >>> X = Frechet("x", a, s, m) + + >>> density(X)(z) + a*((-m + z)/s)**(-a - 1)*exp(-1/((-m + z)/s)**a)/s + + >>> cdf(X)(z) + Piecewise((exp(-1/((-m + z)/s)**a), m <= z), (0, True)) + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Fr%C3%A9chet_distribution + + """ + + return rv(name, FrechetDistribution, (a, s, m)) + +#------------------------------------------------------------------------------- +# Gamma distribution ----------------------------------------------------------- + + +class GammaDistribution(SingleContinuousDistribution): + _argnames = ('k', 'theta') + + set = Interval(0, oo) + + @staticmethod + def check(k, theta): + _value_check(k > 0, "k must be positive") + _value_check(theta > 0, "Theta must be positive") + + def pdf(self, x): + k, theta = self.k, self.theta + return x**(k - 1) * exp(-x/theta) / (gamma(k)*theta**k) + + def _cdf(self, x): + k, theta = self.k, self.theta + return Piecewise( + (lowergamma(k, S(x)/theta)/gamma(k), x > 0), + (S.Zero, True)) + + def _characteristic_function(self, t): + return (1 - self.theta*I*t)**(-self.k) + + def _moment_generating_function(self, t): + return (1- self.theta*t)**(-self.k) + + +def Gamma(name, k, theta): + r""" + Create a continuous random variable with a Gamma distribution. + + Explanation + =========== + + The density of the Gamma distribution is given by + + .. math:: + f(x) := \frac{1}{\Gamma(k) \theta^k} x^{k - 1} e^{-\frac{x}{\theta}} + + with :math:`x \in [0,1]`. + + Parameters + ========== + + k : Real number, `k > 0`, a shape + theta : Real number, `\theta > 0`, a scale + + Returns + ======= + + RandomSymbol + + Examples + ======== + + >>> from sympy.stats import Gamma, density, cdf, E, variance + >>> from sympy import Symbol, pprint, simplify + + >>> k = Symbol("k", positive=True) + >>> theta = Symbol("theta", positive=True) + >>> z = Symbol("z") + + >>> X = Gamma("x", k, theta) + + >>> D = density(X)(z) + >>> pprint(D, use_unicode=False) + -z + ----- + -k k - 1 theta + theta *z *e + --------------------- + Gamma(k) + + >>> C = cdf(X, meijerg=True)(z) + >>> pprint(C, use_unicode=False) + / / z \ + |k*lowergamma|k, -----| + | \ theta/ + <---------------------- for z >= 0 + | Gamma(k + 1) + | + \ 0 otherwise + + >>> E(X) + k*theta + + >>> V = simplify(variance(X)) + >>> pprint(V, use_unicode=False) + 2 + k*theta + + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Gamma_distribution + .. [2] https://mathworld.wolfram.com/GammaDistribution.html + + """ + + return rv(name, GammaDistribution, (k, theta)) + +#------------------------------------------------------------------------------- +# Inverse Gamma distribution --------------------------------------------------- + + +class GammaInverseDistribution(SingleContinuousDistribution): + _argnames = ('a', 'b') + + set = Interval(0, oo) + + @staticmethod + def check(a, b): + _value_check(a > 0, "alpha must be positive") + _value_check(b > 0, "beta must be positive") + + def pdf(self, x): + a, b = self.a, self.b + return b**a/gamma(a) * x**(-a-1) * exp(-b/x) + + def _cdf(self, x): + a, b = self.a, self.b + return Piecewise((uppergamma(a,b/x)/gamma(a), x > 0), + (S.Zero, True)) + + def _characteristic_function(self, t): + a, b = self.a, self.b + return 2 * (-I*b*t)**(a/2) * besselk(a, sqrt(-4*I*b*t)) / gamma(a) + + def _moment_generating_function(self, t): + raise NotImplementedError('The moment generating function for the ' + 'gamma inverse distribution does not exist.') + +def GammaInverse(name, a, b): + r""" + Create a continuous random variable with an inverse Gamma distribution. + + Explanation + =========== + + The density of the inverse Gamma distribution is given by + + .. math:: + f(x) := \frac{\beta^\alpha}{\Gamma(\alpha)} x^{-\alpha - 1} + \exp\left(\frac{-\beta}{x}\right) + + with :math:`x > 0`. + + Parameters + ========== + + a : Real number, `a > 0`, a shape + b : Real number, `b > 0`, a scale + + Returns + ======= + + RandomSymbol + + Examples + ======== + + >>> from sympy.stats import GammaInverse, density, cdf + >>> from sympy import Symbol, pprint + + >>> a = Symbol("a", positive=True) + >>> b = Symbol("b", positive=True) + >>> z = Symbol("z") + + >>> X = GammaInverse("x", a, b) + + >>> D = density(X)(z) + >>> pprint(D, use_unicode=False) + -b + --- + a -a - 1 z + b *z *e + --------------- + Gamma(a) + + >>> cdf(X)(z) + Piecewise((uppergamma(a, b/z)/gamma(a), z > 0), (0, True)) + + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Inverse-gamma_distribution + + """ + + return rv(name, GammaInverseDistribution, (a, b)) + + +#------------------------------------------------------------------------------- +# Gumbel distribution (Maximum and Minimum) -------------------------------------------------------- + + +class GumbelDistribution(SingleContinuousDistribution): + _argnames = ('beta', 'mu', 'minimum') + + set = Interval(-oo, oo) + + @staticmethod + def check(beta, mu, minimum): + _value_check(beta > 0, "Scale parameter beta must be positive.") + + def pdf(self, x): + beta, mu = self.beta, self.mu + z = (x - mu)/beta + f_max = (1/beta)*exp(-z - exp(-z)) + f_min = (1/beta)*exp(z - exp(z)) + return Piecewise((f_min, self.minimum), (f_max, not self.minimum)) + + def _cdf(self, x): + beta, mu = self.beta, self.mu + z = (x - mu)/beta + F_max = exp(-exp(-z)) + F_min = 1 - exp(-exp(z)) + return Piecewise((F_min, self.minimum), (F_max, not self.minimum)) + + def _characteristic_function(self, t): + cf_max = gamma(1 - I*self.beta*t) * exp(I*self.mu*t) + cf_min = gamma(1 + I*self.beta*t) * exp(I*self.mu*t) + return Piecewise((cf_min, self.minimum), (cf_max, not self.minimum)) + + def _moment_generating_function(self, t): + mgf_max = gamma(1 - self.beta*t) * exp(self.mu*t) + mgf_min = gamma(1 + self.beta*t) * exp(self.mu*t) + return Piecewise((mgf_min, self.minimum), (mgf_max, not self.minimum)) + +def Gumbel(name, beta, mu, minimum=False): + r""" + Create a Continuous Random Variable with Gumbel distribution. + + Explanation + =========== + + The density of the Gumbel distribution is given by + + For Maximum + + .. math:: + f(x) := \dfrac{1}{\beta} \exp \left( -\dfrac{x-\mu}{\beta} + - \exp \left( -\dfrac{x - \mu}{\beta} \right) \right) + + with :math:`x \in [ - \infty, \infty ]`. + + For Minimum + + .. math:: + f(x) := \frac{e^{- e^{\frac{- \mu + x}{\beta}} + \frac{- \mu + x}{\beta}}}{\beta} + + with :math:`x \in [ - \infty, \infty ]`. + + Parameters + ========== + + mu : Real number, `\mu`, a location + beta : Real number, `\beta > 0`, a scale + minimum : Boolean, by default ``False``, set to ``True`` for enabling minimum distribution + + Returns + ======= + + RandomSymbol + + Examples + ======== + + >>> from sympy.stats import Gumbel, density, cdf + >>> from sympy import Symbol + >>> x = Symbol("x") + >>> mu = Symbol("mu") + >>> beta = Symbol("beta", positive=True) + >>> X = Gumbel("x", beta, mu) + >>> density(X)(x) + exp(-exp(-(-mu + x)/beta) - (-mu + x)/beta)/beta + >>> cdf(X)(x) + exp(-exp(-(-mu + x)/beta)) + + References + ========== + + .. [1] https://mathworld.wolfram.com/GumbelDistribution.html + .. [2] https://en.wikipedia.org/wiki/Gumbel_distribution + .. [3] https://web.archive.org/web/20200628222206/http://www.mathwave.com/help/easyfit/html/analyses/distributions/gumbel_max.html + .. [4] https://web.archive.org/web/20200628222212/http://www.mathwave.com/help/easyfit/html/analyses/distributions/gumbel_min.html + + """ + return rv(name, GumbelDistribution, (beta, mu, minimum)) + +#------------------------------------------------------------------------------- +# Gompertz distribution -------------------------------------------------------- + +class GompertzDistribution(SingleContinuousDistribution): + _argnames = ('b', 'eta') + + set = Interval(0, oo) + + @staticmethod + def check(b, eta): + _value_check(b > 0, "b must be positive") + _value_check(eta > 0, "eta must be positive") + + def pdf(self, x): + eta, b = self.eta, self.b + return b*eta*exp(b*x)*exp(eta)*exp(-eta*exp(b*x)) + + def _cdf(self, x): + eta, b = self.eta, self.b + return 1 - exp(eta)*exp(-eta*exp(b*x)) + + def _moment_generating_function(self, t): + eta, b = self.eta, self.b + return eta * exp(eta) * expint(t/b, eta) + +def Gompertz(name, b, eta): + r""" + Create a Continuous Random Variable with Gompertz distribution. + + Explanation + =========== + + The density of the Gompertz distribution is given by + + .. math:: + f(x) := b \eta e^{b x} e^{\eta} \exp \left(-\eta e^{bx} \right) + + with :math:`x \in [0, \infty)`. + + Parameters + ========== + + b : Real number, `b > 0`, a scale + eta : Real number, `\eta > 0`, a shape + + Returns + ======= + + RandomSymbol + + Examples + ======== + + >>> from sympy.stats import Gompertz, density + >>> from sympy import Symbol + + >>> b = Symbol("b", positive=True) + >>> eta = Symbol("eta", positive=True) + >>> z = Symbol("z") + + >>> X = Gompertz("x", b, eta) + + >>> density(X)(z) + b*eta*exp(eta)*exp(b*z)*exp(-eta*exp(b*z)) + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Gompertz_distribution + + """ + return rv(name, GompertzDistribution, (b, eta)) + +#------------------------------------------------------------------------------- +# Kumaraswamy distribution ----------------------------------------------------- + + +class KumaraswamyDistribution(SingleContinuousDistribution): + _argnames = ('a', 'b') + + set = Interval(0, oo) + + @staticmethod + def check(a, b): + _value_check(a > 0, "a must be positive") + _value_check(b > 0, "b must be positive") + + def pdf(self, x): + a, b = self.a, self.b + return a * b * x**(a-1) * (1-x**a)**(b-1) + + def _cdf(self, x): + a, b = self.a, self.b + return Piecewise( + (S.Zero, x < S.Zero), + (1 - (1 - x**a)**b, x <= S.One), + (S.One, True)) + +def Kumaraswamy(name, a, b): + r""" + Create a Continuous Random Variable with a Kumaraswamy distribution. + + Explanation + =========== + + The density of the Kumaraswamy distribution is given by + + .. math:: + f(x) := a b x^{a-1} (1-x^a)^{b-1} + + with :math:`x \in [0,1]`. + + Parameters + ========== + + a : Real number, `a > 0`, a shape + b : Real number, `b > 0`, a shape + + Returns + ======= + + RandomSymbol + + Examples + ======== + + >>> from sympy.stats import Kumaraswamy, density, cdf + >>> from sympy import Symbol, pprint + + >>> a = Symbol("a", positive=True) + >>> b = Symbol("b", positive=True) + >>> z = Symbol("z") + + >>> X = Kumaraswamy("x", a, b) + + >>> D = density(X)(z) + >>> pprint(D, use_unicode=False) + b - 1 + a - 1 / a\ + a*b*z *\1 - z / + + >>> cdf(X)(z) + Piecewise((0, z < 0), (1 - (1 - z**a)**b, z <= 1), (1, True)) + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Kumaraswamy_distribution + + """ + + return rv(name, KumaraswamyDistribution, (a, b)) + +#------------------------------------------------------------------------------- +# Laplace distribution --------------------------------------------------------- + + +class LaplaceDistribution(SingleContinuousDistribution): + _argnames = ('mu', 'b') + + set = Interval(-oo, oo) + + @staticmethod + def check(mu, b): + _value_check(b > 0, "Scale parameter b must be positive.") + _value_check(mu.is_real, "Location parameter mu should be real") + + def pdf(self, x): + mu, b = self.mu, self.b + return 1/(2*b)*exp(-Abs(x - mu)/b) + + def _cdf(self, x): + mu, b = self.mu, self.b + return Piecewise( + (S.Half*exp((x - mu)/b), x < mu), + (S.One - S.Half*exp(-(x - mu)/b), x >= mu) + ) + + def _characteristic_function(self, t): + return exp(self.mu*I*t) / (1 + self.b**2*t**2) + + def _moment_generating_function(self, t): + return exp(self.mu*t) / (1 - self.b**2*t**2) + +def Laplace(name, mu, b): + r""" + Create a continuous random variable with a Laplace distribution. + + Explanation + =========== + + The density of the Laplace distribution is given by + + .. math:: + f(x) := \frac{1}{2 b} \exp \left(-\frac{|x-\mu|}b \right) + + Parameters + ========== + + mu : Real number or a list/matrix, the location (mean) or the + location vector + b : Real number or a positive definite matrix, representing a scale + or the covariance matrix. + + Returns + ======= + + RandomSymbol + + Examples + ======== + + >>> from sympy.stats import Laplace, density, cdf + >>> from sympy import Symbol, pprint + + >>> mu = Symbol("mu") + >>> b = Symbol("b", positive=True) + >>> z = Symbol("z") + + >>> X = Laplace("x", mu, b) + + >>> density(X)(z) + exp(-Abs(mu - z)/b)/(2*b) + + >>> cdf(X)(z) + Piecewise((exp((-mu + z)/b)/2, mu > z), (1 - exp((mu - z)/b)/2, True)) + + >>> L = Laplace('L', [1, 2], [[1, 0], [0, 1]]) + >>> pprint(density(L)(1, 2), use_unicode=False) + 5 / ____\ + e *besselk\0, \/ 35 / + --------------------- + pi + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Laplace_distribution + .. [2] https://mathworld.wolfram.com/LaplaceDistribution.html + + """ + + if isinstance(mu, (list, MatrixBase)) and\ + isinstance(b, (list, MatrixBase)): + from sympy.stats.joint_rv_types import MultivariateLaplace + return MultivariateLaplace(name, mu, b) + + return rv(name, LaplaceDistribution, (mu, b)) + +#------------------------------------------------------------------------------- +# Levy distribution --------------------------------------------------------- + + +class LevyDistribution(SingleContinuousDistribution): + _argnames = ('mu', 'c') + + @property + def set(self): + return Interval(self.mu, oo) + + @staticmethod + def check(mu, c): + _value_check(c > 0, "c (scale parameter) must be positive") + _value_check(mu.is_real, "mu (location parameter) must be real") + + def pdf(self, x): + mu, c = self.mu, self.c + return sqrt(c/(2*pi))*exp(-c/(2*(x - mu)))/((x - mu)**(S.One + S.Half)) + + def _cdf(self, x): + mu, c = self.mu, self.c + return erfc(sqrt(c/(2*(x - mu)))) + + def _characteristic_function(self, t): + mu, c = self.mu, self.c + return exp(I * mu * t - sqrt(-2 * I * c * t)) + + def _moment_generating_function(self, t): + raise NotImplementedError('The moment generating function of Levy distribution does not exist.') + +def Levy(name, mu, c): + r""" + Create a continuous random variable with a Levy distribution. + + The density of the Levy distribution is given by + + .. math:: + f(x) := \sqrt(\frac{c}{2 \pi}) \frac{\exp -\frac{c}{2 (x - \mu)}}{(x - \mu)^{3/2}} + + Parameters + ========== + + mu : Real number + The location parameter. + c : Real number, `c > 0` + A scale parameter. + + Returns + ======= + + RandomSymbol + + Examples + ======== + + >>> from sympy.stats import Levy, density, cdf + >>> from sympy import Symbol + + >>> mu = Symbol("mu", real=True) + >>> c = Symbol("c", positive=True) + >>> z = Symbol("z") + + >>> X = Levy("x", mu, c) + + >>> density(X)(z) + sqrt(2)*sqrt(c)*exp(-c/(-2*mu + 2*z))/(2*sqrt(pi)*(-mu + z)**(3/2)) + + >>> cdf(X)(z) + erfc(sqrt(c)*sqrt(1/(-2*mu + 2*z))) + + References + ========== + .. [1] https://en.wikipedia.org/wiki/L%C3%A9vy_distribution + .. [2] https://mathworld.wolfram.com/LevyDistribution.html + """ + + return rv(name, LevyDistribution, (mu, c)) + +#------------------------------------------------------------------------------- +# Log-Cauchy distribution -------------------------------------------------------- + + +class LogCauchyDistribution(SingleContinuousDistribution): + _argnames = ('mu', 'sigma') + + set = Interval.open(0, oo) + + @staticmethod + def check(mu, sigma): + _value_check((sigma > 0) != False, "Scale parameter Gamma must be positive.") + _value_check(mu.is_real != False, "Location parameter must be real.") + + def pdf(self, x): + mu, sigma = self.mu, self.sigma + return 1/(x*pi)*(sigma/((log(x) - mu)**2 + sigma**2)) + + def _cdf(self, x): + mu, sigma = self.mu, self.sigma + return (1/pi)*atan((log(x) - mu)/sigma) + S.Half + + def _characteristic_function(self, t): + raise NotImplementedError("The characteristic function for the " + "Log-Cauchy distribution does not exist.") + + def _moment_generating_function(self, t): + raise NotImplementedError("The moment generating function for the " + "Log-Cauchy distribution does not exist.") + +def LogCauchy(name, mu, sigma): + r""" + Create a continuous random variable with a Log-Cauchy distribution. + The density of the Log-Cauchy distribution is given by + + .. math:: + f(x) := \frac{1}{\pi x} \frac{\sigma}{(log(x)-\mu^2) + \sigma^2} + + Parameters + ========== + + mu : Real number, the location + + sigma : Real number, `\sigma > 0`, a scale + + Returns + ======= + + RandomSymbol + + Examples + ======== + + >>> from sympy.stats import LogCauchy, density, cdf + >>> from sympy import Symbol, S + + >>> mu = 2 + >>> sigma = S.One / 5 + >>> z = Symbol("z") + + >>> X = LogCauchy("x", mu, sigma) + + >>> density(X)(z) + 1/(5*pi*z*((log(z) - 2)**2 + 1/25)) + + >>> cdf(X)(z) + atan(5*log(z) - 10)/pi + 1/2 + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Log-Cauchy_distribution + """ + + return rv(name, LogCauchyDistribution, (mu, sigma)) + + +#------------------------------------------------------------------------------- +# Logistic distribution -------------------------------------------------------- + + +class LogisticDistribution(SingleContinuousDistribution): + _argnames = ('mu', 's') + + set = Interval(-oo, oo) + + @staticmethod + def check(mu, s): + _value_check(s > 0, "Scale parameter s must be positive.") + + def pdf(self, x): + mu, s = self.mu, self.s + return exp(-(x - mu)/s)/(s*(1 + exp(-(x - mu)/s))**2) + + def _cdf(self, x): + mu, s = self.mu, self.s + return S.One/(1 + exp(-(x - mu)/s)) + + def _characteristic_function(self, t): + return Piecewise((exp(I*t*self.mu) * pi*self.s*t / sinh(pi*self.s*t), Ne(t, 0)), (S.One, True)) + + def _moment_generating_function(self, t): + return exp(self.mu*t) * beta_fn(1 - self.s*t, 1 + self.s*t) + + def _quantile(self, p): + return self.mu - self.s*log(-S.One + S.One/p) + +def Logistic(name, mu, s): + r""" + Create a continuous random variable with a logistic distribution. + + Explanation + =========== + + The density of the logistic distribution is given by + + .. math:: + f(x) := \frac{e^{-(x-\mu)/s}} {s\left(1+e^{-(x-\mu)/s}\right)^2} + + Parameters + ========== + + mu : Real number, the location (mean) + s : Real number, `s > 0`, a scale + + Returns + ======= + + RandomSymbol + + Examples + ======== + + >>> from sympy.stats import Logistic, density, cdf + >>> from sympy import Symbol + + >>> mu = Symbol("mu", real=True) + >>> s = Symbol("s", positive=True) + >>> z = Symbol("z") + + >>> X = Logistic("x", mu, s) + + >>> density(X)(z) + exp((mu - z)/s)/(s*(exp((mu - z)/s) + 1)**2) + + >>> cdf(X)(z) + 1/(exp((mu - z)/s) + 1) + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Logistic_distribution + .. [2] https://mathworld.wolfram.com/LogisticDistribution.html + + """ + + return rv(name, LogisticDistribution, (mu, s)) + +#------------------------------------------------------------------------------- +# Log-logistic distribution -------------------------------------------------------- + + +class LogLogisticDistribution(SingleContinuousDistribution): + _argnames = ('alpha', 'beta') + + set = Interval(0, oo) + + @staticmethod + def check(alpha, beta): + _value_check(alpha > 0, "Scale parameter Alpha must be positive.") + _value_check(beta > 0, "Shape parameter Beta must be positive.") + + def pdf(self, x): + a, b = self.alpha, self.beta + return ((b/a)*(x/a)**(b - 1))/(1 + (x/a)**b)**2 + + def _cdf(self, x): + a, b = self.alpha, self.beta + return 1/(1 + (x/a)**(-b)) + + def _quantile(self, p): + a, b = self.alpha, self.beta + return a*((p/(1 - p))**(1/b)) + + def expectation(self, expr, var, **kwargs): + a, b = self.args + return Piecewise((S.NaN, b <= 1), (pi*a/(b*sin(pi/b)), True)) + +def LogLogistic(name, alpha, beta): + r""" + Create a continuous random variable with a log-logistic distribution. + The distribution is unimodal when ``beta > 1``. + + Explanation + =========== + + The density of the log-logistic distribution is given by + + .. math:: + f(x) := \frac{(\frac{\beta}{\alpha})(\frac{x}{\alpha})^{\beta - 1}} + {(1 + (\frac{x}{\alpha})^{\beta})^2} + + Parameters + ========== + + alpha : Real number, `\alpha > 0`, scale parameter and median of distribution + beta : Real number, `\beta > 0`, a shape parameter + + Returns + ======= + + RandomSymbol + + Examples + ======== + + >>> from sympy.stats import LogLogistic, density, cdf, quantile + >>> from sympy import Symbol, pprint + + >>> alpha = Symbol("alpha", positive=True) + >>> beta = Symbol("beta", positive=True) + >>> p = Symbol("p") + >>> z = Symbol("z", positive=True) + + >>> X = LogLogistic("x", alpha, beta) + + >>> D = density(X)(z) + >>> pprint(D, use_unicode=False) + beta - 1 + / z \ + beta*|-----| + \alpha/ + ------------------------ + 2 + / beta \ + |/ z \ | + alpha*||-----| + 1| + \\alpha/ / + + >>> cdf(X)(z) + 1/(1 + (z/alpha)**(-beta)) + + >>> quantile(X)(p) + alpha*(p/(1 - p))**(1/beta) + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Log-logistic_distribution + + """ + + return rv(name, LogLogisticDistribution, (alpha, beta)) + +#------------------------------------------------------------------------------- +#Logit-Normal distribution------------------------------------------------------ + +class LogitNormalDistribution(SingleContinuousDistribution): + _argnames = ('mu', 's') + set = Interval.open(0, 1) + + @staticmethod + def check(mu, s): + _value_check((s ** 2).is_real is not False and s ** 2 > 0, "Squared scale parameter s must be positive.") + _value_check(mu.is_real is not False, "Location parameter must be real") + + def _logit(self, x): + return log(x / (1 - x)) + + def pdf(self, x): + mu, s = self.mu, self.s + return exp(-(self._logit(x) - mu)**2/(2*s**2))*(S.One/sqrt(2*pi*(s**2)))*(1/(x*(1 - x))) + + def _cdf(self, x): + mu, s = self.mu, self.s + return (S.One/2)*(1 + erf((self._logit(x) - mu)/(sqrt(2*s**2)))) + + +def LogitNormal(name, mu, s): + r""" + Create a continuous random variable with a Logit-Normal distribution. + + The density of the logistic distribution is given by + + .. math:: + f(x) := \frac{1}{s \sqrt{2 \pi}} \frac{1}{x(1 - x)} e^{- \frac{(logit(x) - \mu)^2}{s^2}} + where logit(x) = \log(\frac{x}{1 - x}) + Parameters + ========== + + mu : Real number, the location (mean) + s : Real number, `s > 0`, a scale + + Returns + ======= + + RandomSymbol + + Examples + ======== + + >>> from sympy.stats import LogitNormal, density, cdf + >>> from sympy import Symbol,pprint + + >>> mu = Symbol("mu", real=True) + >>> s = Symbol("s", positive=True) + >>> z = Symbol("z") + >>> X = LogitNormal("x",mu,s) + + >>> D = density(X)(z) + >>> pprint(D, use_unicode=False) + 2 + / / z \\ + -|-mu + log|-----|| + \ \1 - z// + --------------------- + 2 + ___ 2*s + \/ 2 *e + ---------------------------- + ____ + 2*\/ pi *s*z*(1 - z) + + >>> density(X)(z) + sqrt(2)*exp(-(-mu + log(z/(1 - z)))**2/(2*s**2))/(2*sqrt(pi)*s*z*(1 - z)) + + >>> cdf(X)(z) + erf(sqrt(2)*(-mu + log(z/(1 - z)))/(2*s))/2 + 1/2 + + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Logit-normal_distribution + + """ + + return rv(name, LogitNormalDistribution, (mu, s)) + +#------------------------------------------------------------------------------- +# Log Normal distribution ------------------------------------------------------ + + +class LogNormalDistribution(SingleContinuousDistribution): + _argnames = ('mean', 'std') + + set = Interval(0, oo) + + @staticmethod + def check(mean, std): + _value_check(std > 0, "Parameter std must be positive.") + + def pdf(self, x): + mean, std = self.mean, self.std + return exp(-(log(x) - mean)**2 / (2*std**2)) / (x*sqrt(2*pi)*std) + + def _cdf(self, x): + mean, std = self.mean, self.std + return Piecewise( + (S.Half + S.Half*erf((log(x) - mean)/sqrt(2)/std), x > 0), + (S.Zero, True) + ) + + def _moment_generating_function(self, t): + raise NotImplementedError('Moment generating function of the log-normal distribution is not defined.') + + +def LogNormal(name, mean, std): + r""" + Create a continuous random variable with a log-normal distribution. + + Explanation + =========== + + The density of the log-normal distribution is given by + + .. math:: + f(x) := \frac{1}{x\sqrt{2\pi\sigma^2}} + e^{-\frac{\left(\ln x-\mu\right)^2}{2\sigma^2}} + + with :math:`x \geq 0`. + + Parameters + ========== + + mu : Real number + The log-scale. + sigma : Real number + A shape. ($\sigma^2 > 0$) + + Returns + ======= + + RandomSymbol + + Examples + ======== + + >>> from sympy.stats import LogNormal, density + >>> from sympy import Symbol, pprint + + >>> mu = Symbol("mu", real=True) + >>> sigma = Symbol("sigma", positive=True) + >>> z = Symbol("z") + + >>> X = LogNormal("x", mu, sigma) + + >>> D = density(X)(z) + >>> pprint(D, use_unicode=False) + 2 + -(-mu + log(z)) + ----------------- + 2 + ___ 2*sigma + \/ 2 *e + ------------------------ + ____ + 2*\/ pi *sigma*z + + + >>> X = LogNormal('x', 0, 1) # Mean 0, standard deviation 1 + + >>> density(X)(z) + sqrt(2)*exp(-log(z)**2/2)/(2*sqrt(pi)*z) + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Lognormal + .. [2] https://mathworld.wolfram.com/LogNormalDistribution.html + + """ + + return rv(name, LogNormalDistribution, (mean, std)) + +#------------------------------------------------------------------------------- +# Lomax Distribution ----------------------------------------------------------- + +class LomaxDistribution(SingleContinuousDistribution): + _argnames = ('alpha', 'lamda',) + set = Interval(0, oo) + + @staticmethod + def check(alpha, lamda): + _value_check(alpha.is_real, "Shape parameter should be real.") + _value_check(lamda.is_real, "Scale parameter should be real.") + _value_check(alpha.is_positive, "Shape parameter should be positive.") + _value_check(lamda.is_positive, "Scale parameter should be positive.") + + def pdf(self, x): + lamba, alpha = self.lamda, self.alpha + return (alpha/lamba) * (S.One + x/lamba)**(-alpha-1) + +def Lomax(name, alpha, lamda): + r""" + Create a continuous random variable with a Lomax distribution. + + Explanation + =========== + + The density of the Lomax distribution is given by + + .. math:: + f(x) := \frac{\alpha}{\lambda}\left[1+\frac{x}{\lambda}\right]^{-(\alpha+1)} + + Parameters + ========== + + alpha : Real Number, `\alpha > 0` + Shape parameter + lamda : Real Number, `\lambda > 0` + Scale parameter + + Examples + ======== + + >>> from sympy.stats import Lomax, density, cdf, E + >>> from sympy import symbols + >>> a, l = symbols('a, l', positive=True) + >>> X = Lomax('X', a, l) + >>> x = symbols('x') + >>> density(X)(x) + a*(1 + x/l)**(-a - 1)/l + >>> cdf(X)(x) + Piecewise((1 - 1/(1 + x/l)**a, x >= 0), (0, True)) + >>> a = 2 + >>> X = Lomax('X', a, l) + >>> E(X) + l + + Returns + ======= + + RandomSymbol + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Lomax_distribution + + """ + return rv(name, LomaxDistribution, (alpha, lamda)) + +#------------------------------------------------------------------------------- +# Maxwell distribution --------------------------------------------------------- + + +class MaxwellDistribution(SingleContinuousDistribution): + _argnames = ('a',) + + set = Interval(0, oo) + + @staticmethod + def check(a): + _value_check(a > 0, "Parameter a must be positive.") + + def pdf(self, x): + a = self.a + return sqrt(2/pi)*x**2*exp(-x**2/(2*a**2))/a**3 + + def _cdf(self, x): + a = self.a + return erf(sqrt(2)*x/(2*a)) - sqrt(2)*x*exp(-x**2/(2*a**2))/(sqrt(pi)*a) + +def Maxwell(name, a): + r""" + Create a continuous random variable with a Maxwell distribution. + + Explanation + =========== + + The density of the Maxwell distribution is given by + + .. math:: + f(x) := \sqrt{\frac{2}{\pi}} \frac{x^2 e^{-x^2/(2a^2)}}{a^3} + + with :math:`x \geq 0`. + + .. TODO - what does the parameter mean? + + Parameters + ========== + + a : Real number, `a > 0` + + Returns + ======= + + RandomSymbol + + Examples + ======== + + >>> from sympy.stats import Maxwell, density, E, variance + >>> from sympy import Symbol, simplify + + >>> a = Symbol("a", positive=True) + >>> z = Symbol("z") + + >>> X = Maxwell("x", a) + + >>> density(X)(z) + sqrt(2)*z**2*exp(-z**2/(2*a**2))/(sqrt(pi)*a**3) + + >>> E(X) + 2*sqrt(2)*a/sqrt(pi) + + >>> simplify(variance(X)) + a**2*(-8 + 3*pi)/pi + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Maxwell_distribution + .. [2] https://mathworld.wolfram.com/MaxwellDistribution.html + + """ + + return rv(name, MaxwellDistribution, (a, )) + +#------------------------------------------------------------------------------- +# Moyal Distribution ----------------------------------------------------------- +class MoyalDistribution(SingleContinuousDistribution): + _argnames = ('mu', 'sigma') + + @staticmethod + def check(mu, sigma): + _value_check(mu.is_real, "Location parameter must be real.") + _value_check(sigma.is_real and sigma > 0, "Scale parameter must be real\ + and positive.") + + def pdf(self, x): + mu, sigma = self.mu, self.sigma + num = exp(-(exp(-(x - mu)/sigma) + (x - mu)/(sigma))/2) + den = (sqrt(2*pi) * sigma) + return num/den + + def _characteristic_function(self, t): + mu, sigma = self.mu, self.sigma + term1 = exp(I*t*mu) + term2 = (2**(-I*sigma*t) * gamma(Rational(1, 2) - I*t*sigma)) + return (term1 * term2)/sqrt(pi) + + def _moment_generating_function(self, t): + mu, sigma = self.mu, self.sigma + term1 = exp(t*mu) + term2 = (2**(-1*sigma*t) * gamma(Rational(1, 2) - t*sigma)) + return (term1 * term2)/sqrt(pi) + +def Moyal(name, mu, sigma): + r""" + Create a continuous random variable with a Moyal distribution. + + Explanation + =========== + + The density of the Moyal distribution is given by + + .. math:: + f(x) := \frac{\exp-\frac{1}{2}\exp-\frac{x-\mu}{\sigma}-\frac{x-\mu}{2\sigma}}{\sqrt{2\pi}\sigma} + + with :math:`x \in \mathbb{R}`. + + Parameters + ========== + + mu : Real number + Location parameter + sigma : Real positive number + Scale parameter + + Returns + ======= + + RandomSymbol + + Examples + ======== + + >>> from sympy.stats import Moyal, density, cdf + >>> from sympy import Symbol, simplify + >>> mu = Symbol("mu", real=True) + >>> sigma = Symbol("sigma", positive=True, real=True) + >>> z = Symbol("z") + >>> X = Moyal("x", mu, sigma) + >>> density(X)(z) + sqrt(2)*exp(-exp((mu - z)/sigma)/2 - (-mu + z)/(2*sigma))/(2*sqrt(pi)*sigma) + >>> simplify(cdf(X)(z)) + 1 - erf(sqrt(2)*exp((mu - z)/(2*sigma))/2) + + References + ========== + + .. [1] https://reference.wolfram.com/language/ref/MoyalDistribution.html + .. [2] https://www.stat.rice.edu/~dobelman/textfiles/DistributionsHandbook.pdf + + """ + + return rv(name, MoyalDistribution, (mu, sigma)) + +#------------------------------------------------------------------------------- +# Nakagami distribution -------------------------------------------------------- + + +class NakagamiDistribution(SingleContinuousDistribution): + _argnames = ('mu', 'omega') + + set = Interval(0, oo) + + @staticmethod + def check(mu, omega): + _value_check(mu >= S.Half, "Shape parameter mu must be greater than equal to 1/2.") + _value_check(omega > 0, "Spread parameter omega must be positive.") + + def pdf(self, x): + mu, omega = self.mu, self.omega + return 2*mu**mu/(gamma(mu)*omega**mu)*x**(2*mu - 1)*exp(-mu/omega*x**2) + + def _cdf(self, x): + mu, omega = self.mu, self.omega + return Piecewise( + (lowergamma(mu, (mu/omega)*x**2)/gamma(mu), x > 0), + (S.Zero, True)) + +def Nakagami(name, mu, omega): + r""" + Create a continuous random variable with a Nakagami distribution. + + Explanation + =========== + + The density of the Nakagami distribution is given by + + .. math:: + f(x) := \frac{2\mu^\mu}{\Gamma(\mu)\omega^\mu} x^{2\mu-1} + \exp\left(-\frac{\mu}{\omega}x^2 \right) + + with :math:`x > 0`. + + Parameters + ========== + + mu : Real number, `\mu \geq \frac{1}{2}`, a shape + omega : Real number, `\omega > 0`, the spread + + Returns + ======= + + RandomSymbol + + Examples + ======== + + >>> from sympy.stats import Nakagami, density, E, variance, cdf + >>> from sympy import Symbol, simplify, pprint + + >>> mu = Symbol("mu", positive=True) + >>> omega = Symbol("omega", positive=True) + >>> z = Symbol("z") + + >>> X = Nakagami("x", mu, omega) + + >>> D = density(X)(z) + >>> pprint(D, use_unicode=False) + 2 + -mu*z + ------- + mu -mu 2*mu - 1 omega + 2*mu *omega *z *e + ---------------------------------- + Gamma(mu) + + >>> simplify(E(X)) + sqrt(mu)*sqrt(omega)*gamma(mu + 1/2)/gamma(mu + 1) + + >>> V = simplify(variance(X)) + >>> pprint(V, use_unicode=False) + 2 + omega*Gamma (mu + 1/2) + omega - ----------------------- + Gamma(mu)*Gamma(mu + 1) + + >>> cdf(X)(z) + Piecewise((lowergamma(mu, mu*z**2/omega)/gamma(mu), z > 0), + (0, True)) + + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Nakagami_distribution + + """ + + return rv(name, NakagamiDistribution, (mu, omega)) + +#------------------------------------------------------------------------------- +# Normal distribution ---------------------------------------------------------- + + +class NormalDistribution(SingleContinuousDistribution): + _argnames = ('mean', 'std') + + @staticmethod + def check(mean, std): + _value_check(std > 0, "Standard deviation must be positive") + + def pdf(self, x): + return exp(-(x - self.mean)**2 / (2*self.std**2)) / (sqrt(2*pi)*self.std) + + def _cdf(self, x): + mean, std = self.mean, self.std + return erf(sqrt(2)*(-mean + x)/(2*std))/2 + S.Half + + def _characteristic_function(self, t): + mean, std = self.mean, self.std + return exp(I*mean*t - std**2*t**2/2) + + def _moment_generating_function(self, t): + mean, std = self.mean, self.std + return exp(mean*t + std**2*t**2/2) + + def _quantile(self, p): + mean, std = self.mean, self.std + return mean + std*sqrt(2)*erfinv(2*p - 1) + + +def Normal(name, mean, std): + r""" + Create a continuous random variable with a Normal distribution. + + Explanation + =========== + + The density of the Normal distribution is given by + + .. math:: + f(x) := \frac{1}{\sigma\sqrt{2\pi}} e^{ -\frac{(x-\mu)^2}{2\sigma^2} } + + Parameters + ========== + + mu : Real number or a list representing the mean or the mean vector + sigma : Real number or a positive definite square matrix, + :math:`\sigma^2 > 0`, the variance + + Returns + ======= + + RandomSymbol + + Examples + ======== + + >>> from sympy.stats import Normal, density, E, std, cdf, skewness, quantile, marginal_distribution + >>> from sympy import Symbol, simplify, pprint + + >>> mu = Symbol("mu") + >>> sigma = Symbol("sigma", positive=True) + >>> z = Symbol("z") + >>> y = Symbol("y") + >>> p = Symbol("p") + >>> X = Normal("x", mu, sigma) + + >>> density(X)(z) + sqrt(2)*exp(-(-mu + z)**2/(2*sigma**2))/(2*sqrt(pi)*sigma) + + >>> C = simplify(cdf(X))(z) # it needs a little more help... + >>> pprint(C, use_unicode=False) + / ___ \ + |\/ 2 *(-mu + z)| + erf|---------------| + \ 2*sigma / 1 + -------------------- + - + 2 2 + + >>> quantile(X)(p) + mu + sqrt(2)*sigma*erfinv(2*p - 1) + + >>> simplify(skewness(X)) + 0 + + >>> X = Normal("x", 0, 1) # Mean 0, standard deviation 1 + >>> density(X)(z) + sqrt(2)*exp(-z**2/2)/(2*sqrt(pi)) + + >>> E(2*X + 1) + 1 + + >>> simplify(std(2*X + 1)) + 2 + + >>> m = Normal('X', [1, 2], [[2, 1], [1, 2]]) + >>> pprint(density(m)(y, z), use_unicode=False) + 2 2 + y y*z z + - -- + --- - -- + z - 1 + ___ 3 3 3 + \/ 3 *e + ------------------------------ + 6*pi + + >>> marginal_distribution(m, m[0])(1) + 1/(2*sqrt(pi)) + + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Normal_distribution + .. [2] https://mathworld.wolfram.com/NormalDistributionFunction.html + + """ + + if isinstance(mean, list) or getattr(mean, 'is_Matrix', False) and\ + isinstance(std, list) or getattr(std, 'is_Matrix', False): + from sympy.stats.joint_rv_types import MultivariateNormal + return MultivariateNormal(name, mean, std) + return rv(name, NormalDistribution, (mean, std)) + + +#------------------------------------------------------------------------------- +# Inverse Gaussian distribution ---------------------------------------------------------- + + +class GaussianInverseDistribution(SingleContinuousDistribution): + _argnames = ('mean', 'shape') + + @property + def set(self): + return Interval(0, oo) + + @staticmethod + def check(mean, shape): + _value_check(shape > 0, "Shape parameter must be positive") + _value_check(mean > 0, "Mean must be positive") + + def pdf(self, x): + mu, s = self.mean, self.shape + return exp(-s*(x - mu)**2 / (2*x*mu**2)) * sqrt(s/(2*pi*x**3)) + + def _cdf(self, x): + from sympy.stats import cdf + mu, s = self.mean, self.shape + stdNormalcdf = cdf(Normal('x', 0, 1)) + + first_term = stdNormalcdf(sqrt(s/x) * ((x/mu) - S.One)) + second_term = exp(2*s/mu) * stdNormalcdf(-sqrt(s/x)*(x/mu + S.One)) + + return first_term + second_term + + def _characteristic_function(self, t): + mu, s = self.mean, self.shape + return exp((s/mu)*(1 - sqrt(1 - (2*mu**2*I*t)/s))) + + def _moment_generating_function(self, t): + mu, s = self.mean, self.shape + return exp((s/mu)*(1 - sqrt(1 - (2*mu**2*t)/s))) + + +def GaussianInverse(name, mean, shape): + r""" + Create a continuous random variable with an Inverse Gaussian distribution. + Inverse Gaussian distribution is also known as Wald distribution. + + Explanation + =========== + + The density of the Inverse Gaussian distribution is given by + + .. math:: + f(x) := \sqrt{\frac{\lambda}{2\pi x^3}} e^{-\frac{\lambda(x-\mu)^2}{2x\mu^2}} + + Parameters + ========== + + mu : + Positive number representing the mean. + lambda : + Positive number representing the shape parameter. + + Returns + ======= + + RandomSymbol + + Examples + ======== + + >>> from sympy.stats import GaussianInverse, density, E, std, skewness + >>> from sympy import Symbol, pprint + + >>> mu = Symbol("mu", positive=True) + >>> lamda = Symbol("lambda", positive=True) + >>> z = Symbol("z", positive=True) + >>> X = GaussianInverse("x", mu, lamda) + + >>> D = density(X)(z) + >>> pprint(D, use_unicode=False) + 2 + -lambda*(-mu + z) + ------------------- + 2 + ___ ________ 2*mu *z + \/ 2 *\/ lambda *e + ------------------------------------- + ____ 3/2 + 2*\/ pi *z + + >>> E(X) + mu + + >>> std(X).expand() + mu**(3/2)/sqrt(lambda) + + >>> skewness(X).expand() + 3*sqrt(mu)/sqrt(lambda) + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Inverse_Gaussian_distribution + .. [2] https://mathworld.wolfram.com/InverseGaussianDistribution.html + + """ + + return rv(name, GaussianInverseDistribution, (mean, shape)) + +Wald = GaussianInverse + +#------------------------------------------------------------------------------- +# Pareto distribution ---------------------------------------------------------- + + +class ParetoDistribution(SingleContinuousDistribution): + _argnames = ('xm', 'alpha') + + @property + def set(self): + return Interval(self.xm, oo) + + @staticmethod + def check(xm, alpha): + _value_check(xm > 0, "Xm must be positive") + _value_check(alpha > 0, "Alpha must be positive") + + def pdf(self, x): + xm, alpha = self.xm, self.alpha + return alpha * xm**alpha / x**(alpha + 1) + + def _cdf(self, x): + xm, alpha = self.xm, self.alpha + return Piecewise( + (S.One - xm**alpha/x**alpha, x>=xm), + (0, True), + ) + + def _moment_generating_function(self, t): + xm, alpha = self.xm, self.alpha + return alpha * (-xm*t)**alpha * uppergamma(-alpha, -xm*t) + + def _characteristic_function(self, t): + xm, alpha = self.xm, self.alpha + return alpha * (-I * xm * t) ** alpha * uppergamma(-alpha, -I * xm * t) + + +def Pareto(name, xm, alpha): + r""" + Create a continuous random variable with the Pareto distribution. + + Explanation + =========== + + The density of the Pareto distribution is given by + + .. math:: + f(x) := \frac{\alpha\,x_m^\alpha}{x^{\alpha+1}} + + with :math:`x \in [x_m,\infty]`. + + Parameters + ========== + + xm : Real number, `x_m > 0`, a scale + alpha : Real number, `\alpha > 0`, a shape + + Returns + ======= + + RandomSymbol + + Examples + ======== + + >>> from sympy.stats import Pareto, density + >>> from sympy import Symbol + + >>> xm = Symbol("xm", positive=True) + >>> beta = Symbol("beta", positive=True) + >>> z = Symbol("z") + + >>> X = Pareto("x", xm, beta) + + >>> density(X)(z) + beta*xm**beta*z**(-beta - 1) + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Pareto_distribution + .. [2] https://mathworld.wolfram.com/ParetoDistribution.html + + """ + + return rv(name, ParetoDistribution, (xm, alpha)) + +#------------------------------------------------------------------------------- +# PowerFunction distribution --------------------------------------------------- + + +class PowerFunctionDistribution(SingleContinuousDistribution): + _argnames=('alpha','a','b') + + @property + def set(self): + return Interval(self.a, self.b) + + @staticmethod + def check(alpha, a, b): + _value_check(a.is_real, "Continuous Boundary parameter should be real.") + _value_check(b.is_real, "Continuous Boundary parameter should be real.") + _value_check(a < b, " 'a' the left Boundary must be smaller than 'b' the right Boundary." ) + _value_check(alpha.is_positive, "Continuous Shape parameter should be positive.") + + def pdf(self, x): + alpha, a, b = self.alpha, self.a, self.b + num = alpha*(x - a)**(alpha - 1) + den = (b - a)**alpha + return num/den + +def PowerFunction(name, alpha, a, b): + r""" + Creates a continuous random variable with a Power Function Distribution. + + Explanation + =========== + + The density of PowerFunction distribution is given by + + .. math:: + f(x) := \frac{{\alpha}(x - a)^{\alpha - 1}}{(b - a)^{\alpha}} + + with :math:`x \in [a,b]`. + + Parameters + ========== + + alpha : Positive number, `0 < \alpha`, the shape parameter + a : Real number, :math:`-\infty < a`, the left boundary + b : Real number, :math:`a < b < \infty`, the right boundary + + Returns + ======= + + RandomSymbol + + Examples + ======== + + >>> from sympy.stats import PowerFunction, density, cdf, E, variance + >>> from sympy import Symbol + >>> alpha = Symbol("alpha", positive=True) + >>> a = Symbol("a", real=True) + >>> b = Symbol("b", real=True) + >>> z = Symbol("z") + + >>> X = PowerFunction("X", 2, a, b) + + >>> density(X)(z) + (-2*a + 2*z)/(-a + b)**2 + + >>> cdf(X)(z) + Piecewise((a**2/(a**2 - 2*a*b + b**2) - 2*a*z/(a**2 - 2*a*b + b**2) + + z**2/(a**2 - 2*a*b + b**2), a <= z), (0, True)) + + >>> alpha = 2 + >>> a = 0 + >>> b = 1 + >>> Y = PowerFunction("Y", alpha, a, b) + + >>> E(Y) + 2/3 + + >>> variance(Y) + 1/18 + + References + ========== + + .. [1] https://web.archive.org/web/20200204081320/http://www.mathwave.com/help/easyfit/html/analyses/distributions/power_func.html + + """ + return rv(name, PowerFunctionDistribution, (alpha, a, b)) + +#------------------------------------------------------------------------------- +# QuadraticU distribution ------------------------------------------------------ + + +class QuadraticUDistribution(SingleContinuousDistribution): + _argnames = ('a', 'b') + + @property + def set(self): + return Interval(self.a, self.b) + + @staticmethod + def check(a, b): + _value_check(b > a, "Parameter b must be in range (%s, oo)."%(a)) + + def pdf(self, x): + a, b = self.a, self.b + alpha = 12 / (b-a)**3 + beta = (a+b) / 2 + return Piecewise( + (alpha * (x-beta)**2, And(a<=x, x<=b)), + (S.Zero, True)) + + def _moment_generating_function(self, t): + a, b = self.a, self.b + return -3 * (exp(a*t) * (4 + (a**2 + 2*a*(-2 + b) + b**2) * t) \ + - exp(b*t) * (4 + (-4*b + (a + b)**2) * t)) / ((a-b)**3 * t**2) + + def _characteristic_function(self, t): + a, b = self.a, self.b + return -3*I*(exp(I*a*t*exp(I*b*t)) * (4*I - (-4*b + (a+b)**2)*t)) \ + / ((a-b)**3 * t**2) + + +def QuadraticU(name, a, b): + r""" + Create a Continuous Random Variable with a U-quadratic distribution. + + Explanation + =========== + + The density of the U-quadratic distribution is given by + + .. math:: + f(x) := \alpha (x-\beta)^2 + + with :math:`x \in [a,b]`. + + Parameters + ========== + + a : Real number + b : Real number, :math:`a < b` + + Returns + ======= + + RandomSymbol + + Examples + ======== + + >>> from sympy.stats import QuadraticU, density + >>> from sympy import Symbol, pprint + + >>> a = Symbol("a", real=True) + >>> b = Symbol("b", real=True) + >>> z = Symbol("z") + + >>> X = QuadraticU("x", a, b) + + >>> D = density(X)(z) + >>> pprint(D, use_unicode=False) + / 2 + | / a b \ + |12*|- - - - + z| + | \ 2 2 / + <----------------- for And(b >= z, a <= z) + | 3 + | (-a + b) + | + \ 0 otherwise + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/U-quadratic_distribution + + """ + + return rv(name, QuadraticUDistribution, (a, b)) + +#------------------------------------------------------------------------------- +# RaisedCosine distribution ---------------------------------------------------- + + +class RaisedCosineDistribution(SingleContinuousDistribution): + _argnames = ('mu', 's') + + @property + def set(self): + return Interval(self.mu - self.s, self.mu + self.s) + + @staticmethod + def check(mu, s): + _value_check(s > 0, "s must be positive") + + def pdf(self, x): + mu, s = self.mu, self.s + return Piecewise( + ((1+cos(pi*(x-mu)/s)) / (2*s), And(mu-s<=x, x<=mu+s)), + (S.Zero, True)) + + def _characteristic_function(self, t): + mu, s = self.mu, self.s + return Piecewise((exp(-I*pi*mu/s)/2, Eq(t, -pi/s)), + (exp(I*pi*mu/s)/2, Eq(t, pi/s)), + (pi**2*sin(s*t)*exp(I*mu*t) / (s*t*(pi**2 - s**2*t**2)), True)) + + def _moment_generating_function(self, t): + mu, s = self.mu, self.s + return pi**2 * sinh(s*t) * exp(mu*t) / (s*t*(pi**2 + s**2*t**2)) + +def RaisedCosine(name, mu, s): + r""" + Create a Continuous Random Variable with a raised cosine distribution. + + Explanation + =========== + + The density of the raised cosine distribution is given by + + .. math:: + f(x) := \frac{1}{2s}\left(1+\cos\left(\frac{x-\mu}{s}\pi\right)\right) + + with :math:`x \in [\mu-s,\mu+s]`. + + Parameters + ========== + + mu : Real number + s : Real number, `s > 0` + + Returns + ======= + + RandomSymbol + + Examples + ======== + + >>> from sympy.stats import RaisedCosine, density + >>> from sympy import Symbol, pprint + + >>> mu = Symbol("mu", real=True) + >>> s = Symbol("s", positive=True) + >>> z = Symbol("z") + + >>> X = RaisedCosine("x", mu, s) + + >>> D = density(X)(z) + >>> pprint(D, use_unicode=False) + / /pi*(-mu + z)\ + |cos|------------| + 1 + | \ s / + <--------------------- for And(z >= mu - s, z <= mu + s) + | 2*s + | + \ 0 otherwise + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Raised_cosine_distribution + + """ + + return rv(name, RaisedCosineDistribution, (mu, s)) + +#------------------------------------------------------------------------------- +# Rayleigh distribution -------------------------------------------------------- + + +class RayleighDistribution(SingleContinuousDistribution): + _argnames = ('sigma',) + + set = Interval(0, oo) + + @staticmethod + def check(sigma): + _value_check(sigma > 0, "Scale parameter sigma must be positive.") + + def pdf(self, x): + sigma = self.sigma + return x/sigma**2*exp(-x**2/(2*sigma**2)) + + def _cdf(self, x): + sigma = self.sigma + return 1 - exp(-(x**2/(2*sigma**2))) + + def _characteristic_function(self, t): + sigma = self.sigma + return 1 - sigma*t*exp(-sigma**2*t**2/2) * sqrt(pi/2) * (erfi(sigma*t/sqrt(2)) - I) + + def _moment_generating_function(self, t): + sigma = self.sigma + return 1 + sigma*t*exp(sigma**2*t**2/2) * sqrt(pi/2) * (erf(sigma*t/sqrt(2)) + 1) + + +def Rayleigh(name, sigma): + r""" + Create a continuous random variable with a Rayleigh distribution. + + Explanation + =========== + + The density of the Rayleigh distribution is given by + + .. math :: + f(x) := \frac{x}{\sigma^2} e^{-x^2/2\sigma^2} + + with :math:`x > 0`. + + Parameters + ========== + + sigma : Real number, `\sigma > 0` + + Returns + ======= + + RandomSymbol + + Examples + ======== + + >>> from sympy.stats import Rayleigh, density, E, variance + >>> from sympy import Symbol + + >>> sigma = Symbol("sigma", positive=True) + >>> z = Symbol("z") + + >>> X = Rayleigh("x", sigma) + + >>> density(X)(z) + z*exp(-z**2/(2*sigma**2))/sigma**2 + + >>> E(X) + sqrt(2)*sqrt(pi)*sigma/2 + + >>> variance(X) + -pi*sigma**2/2 + 2*sigma**2 + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Rayleigh_distribution + .. [2] https://mathworld.wolfram.com/RayleighDistribution.html + + """ + + return rv(name, RayleighDistribution, (sigma, )) + +#------------------------------------------------------------------------------- +# Reciprocal distribution -------------------------------------------------------- + +class ReciprocalDistribution(SingleContinuousDistribution): + _argnames = ('a', 'b') + + @property + def set(self): + return Interval(self.a, self.b) + + @staticmethod + def check(a, b): + _value_check(a > 0, "Parameter > 0. a = %s"%a) + _value_check((a < b), + "Parameter b must be in range (%s, +oo]. b = %s"%(a, b)) + + def pdf(self, x): + a, b = self.a, self.b + return 1/(x*(log(b) - log(a))) + + +def Reciprocal(name, a, b): + r"""Creates a continuous random variable with a reciprocal distribution. + + + Parameters + ========== + + a : Real number, :math:`0 < a` + b : Real number, :math:`a < b` + + Returns + ======= + + RandomSymbol + + Examples + ======== + + >>> from sympy.stats import Reciprocal, density, cdf + >>> from sympy import symbols + >>> a, b, x = symbols('a, b, x', positive=True) + >>> R = Reciprocal('R', a, b) + + >>> density(R)(x) + 1/(x*(-log(a) + log(b))) + >>> cdf(R)(x) + Piecewise((log(a)/(log(a) - log(b)) - log(x)/(log(a) - log(b)), a <= x), (0, True)) + + Reference + ========= + + .. [1] https://en.wikipedia.org/wiki/Reciprocal_distribution + + """ + return rv(name, ReciprocalDistribution, (a, b)) + + +#------------------------------------------------------------------------------- +# Shifted Gompertz distribution ------------------------------------------------ + + +class ShiftedGompertzDistribution(SingleContinuousDistribution): + _argnames = ('b', 'eta') + + set = Interval(0, oo) + + @staticmethod + def check(b, eta): + _value_check(b > 0, "b must be positive") + _value_check(eta > 0, "eta must be positive") + + def pdf(self, x): + b, eta = self.b, self.eta + return b*exp(-b*x)*exp(-eta*exp(-b*x))*(1+eta*(1-exp(-b*x))) + +def ShiftedGompertz(name, b, eta): + r""" + Create a continuous random variable with a Shifted Gompertz distribution. + + Explanation + =========== + + The density of the Shifted Gompertz distribution is given by + + .. math:: + f(x) := b e^{-b x} e^{-\eta \exp(-b x)} \left[1 + \eta(1 - e^(-bx)) \right] + + with :math:`x \in [0, \infty)`. + + Parameters + ========== + + b : Real number, `b > 0`, a scale + eta : Real number, `\eta > 0`, a shape + + Returns + ======= + + RandomSymbol + + Examples + ======== + >>> from sympy.stats import ShiftedGompertz, density + >>> from sympy import Symbol + + >>> b = Symbol("b", positive=True) + >>> eta = Symbol("eta", positive=True) + >>> x = Symbol("x") + + >>> X = ShiftedGompertz("x", b, eta) + + >>> density(X)(x) + b*(eta*(1 - exp(-b*x)) + 1)*exp(-b*x)*exp(-eta*exp(-b*x)) + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Shifted_Gompertz_distribution + + """ + return rv(name, ShiftedGompertzDistribution, (b, eta)) + +#------------------------------------------------------------------------------- +# StudentT distribution -------------------------------------------------------- + + +class StudentTDistribution(SingleContinuousDistribution): + _argnames = ('nu',) + + set = Interval(-oo, oo) + + @staticmethod + def check(nu): + _value_check(nu > 0, "Degrees of freedom nu must be positive.") + + def pdf(self, x): + nu = self.nu + return 1/(sqrt(nu)*beta_fn(S.Half, nu/2))*(1 + x**2/nu)**(-(nu + 1)/2) + + def _cdf(self, x): + nu = self.nu + return S.Half + x*gamma((nu+1)/2)*hyper((S.Half, (nu+1)/2), + (Rational(3, 2),), -x**2/nu)/(sqrt(pi*nu)*gamma(nu/2)) + + def _moment_generating_function(self, t): + raise NotImplementedError('The moment generating function for the Student-T distribution is undefined.') + + +def StudentT(name, nu): + r""" + Create a continuous random variable with a student's t distribution. + + Explanation + =========== + + The density of the student's t distribution is given by + + .. math:: + f(x) := \frac{\Gamma \left(\frac{\nu+1}{2} \right)} + {\sqrt{\nu\pi}\Gamma \left(\frac{\nu}{2} \right)} + \left(1+\frac{x^2}{\nu} \right)^{-\frac{\nu+1}{2}} + + Parameters + ========== + + nu : Real number, `\nu > 0`, the degrees of freedom + + Returns + ======= + + RandomSymbol + + Examples + ======== + + >>> from sympy.stats import StudentT, density, cdf + >>> from sympy import Symbol, pprint + + >>> nu = Symbol("nu", positive=True) + >>> z = Symbol("z") + + >>> X = StudentT("x", nu) + + >>> D = density(X)(z) + >>> pprint(D, use_unicode=False) + nu 1 + - -- - - + 2 2 + / 2\ + | z | + |1 + --| + \ nu/ + ----------------- + ____ / nu\ + \/ nu *B|1/2, --| + \ 2 / + + >>> cdf(X)(z) + 1/2 + z*gamma(nu/2 + 1/2)*hyper((1/2, nu/2 + 1/2), (3/2,), + -z**2/nu)/(sqrt(pi)*sqrt(nu)*gamma(nu/2)) + + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Student_t-distribution + .. [2] https://mathworld.wolfram.com/Studentst-Distribution.html + + """ + + return rv(name, StudentTDistribution, (nu, )) + +#------------------------------------------------------------------------------- +# Trapezoidal distribution ------------------------------------------------------ + + +class TrapezoidalDistribution(SingleContinuousDistribution): + _argnames = ('a', 'b', 'c', 'd') + + @property + def set(self): + return Interval(self.a, self.d) + + @staticmethod + def check(a, b, c, d): + _value_check(a < d, "Lower bound parameter a < %s. a = %s"%(d, a)) + _value_check((a <= b, b < c), + "Level start parameter b must be in range [%s, %s). b = %s"%(a, c, b)) + _value_check((b < c, c <= d), + "Level end parameter c must be in range (%s, %s]. c = %s"%(b, d, c)) + _value_check(d >= c, "Upper bound parameter d > %s. d = %s"%(c, d)) + + def pdf(self, x): + a, b, c, d = self.a, self.b, self.c, self.d + return Piecewise( + (2*(x-a) / ((b-a)*(d+c-a-b)), And(a <= x, x < b)), + (2 / (d+c-a-b), And(b <= x, x < c)), + (2*(d-x) / ((d-c)*(d+c-a-b)), And(c <= x, x <= d)), + (S.Zero, True)) + +def Trapezoidal(name, a, b, c, d): + r""" + Create a continuous random variable with a trapezoidal distribution. + + Explanation + =========== + + The density of the trapezoidal distribution is given by + + .. math:: + f(x) := \begin{cases} + 0 & \mathrm{for\ } x < a, \\ + \frac{2(x-a)}{(b-a)(d+c-a-b)} & \mathrm{for\ } a \le x < b, \\ + \frac{2}{d+c-a-b} & \mathrm{for\ } b \le x < c, \\ + \frac{2(d-x)}{(d-c)(d+c-a-b)} & \mathrm{for\ } c \le x < d, \\ + 0 & \mathrm{for\ } d < x. + \end{cases} + + Parameters + ========== + + a : Real number, :math:`a < d` + b : Real number, :math:`a \le b < c` + c : Real number, :math:`b < c \le d` + d : Real number + + Returns + ======= + + RandomSymbol + + Examples + ======== + + >>> from sympy.stats import Trapezoidal, density + >>> from sympy import Symbol, pprint + + >>> a = Symbol("a") + >>> b = Symbol("b") + >>> c = Symbol("c") + >>> d = Symbol("d") + >>> z = Symbol("z") + + >>> X = Trapezoidal("x", a,b,c,d) + + >>> pprint(density(X)(z), use_unicode=False) + / -2*a + 2*z + |------------------------- for And(a <= z, b > z) + |(-a + b)*(-a - b + c + d) + | + | 2 + | -------------- for And(b <= z, c > z) + < -a - b + c + d + | + | 2*d - 2*z + |------------------------- for And(d >= z, c <= z) + |(-c + d)*(-a - b + c + d) + | + \ 0 otherwise + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Trapezoidal_distribution + + """ + return rv(name, TrapezoidalDistribution, (a, b, c, d)) + +#------------------------------------------------------------------------------- +# Triangular distribution ------------------------------------------------------ + + +class TriangularDistribution(SingleContinuousDistribution): + _argnames = ('a', 'b', 'c') + + @property + def set(self): + return Interval(self.a, self.b) + + @staticmethod + def check(a, b, c): + _value_check(b > a, "Parameter b > %s. b = %s"%(a, b)) + _value_check((a <= c, c <= b), + "Parameter c must be in range [%s, %s]. c = %s"%(a, b, c)) + + def pdf(self, x): + a, b, c = self.a, self.b, self.c + return Piecewise( + (2*(x - a)/((b - a)*(c - a)), And(a <= x, x < c)), + (2/(b - a), Eq(x, c)), + (2*(b - x)/((b - a)*(b - c)), And(c < x, x <= b)), + (S.Zero, True)) + + def _characteristic_function(self, t): + a, b, c = self.a, self.b, self.c + return -2 *((b-c) * exp(I*a*t) - (b-a) * exp(I*c*t) + (c-a) * exp(I*b*t)) / ((b-a)*(c-a)*(b-c)*t**2) + + def _moment_generating_function(self, t): + a, b, c = self.a, self.b, self.c + return 2 * ((b - c) * exp(a * t) - (b - a) * exp(c * t) + (c - a) * exp(b * t)) / ( + (b - a) * (c - a) * (b - c) * t ** 2) + + +def Triangular(name, a, b, c): + r""" + Create a continuous random variable with a triangular distribution. + + Explanation + =========== + + The density of the triangular distribution is given by + + .. math:: + f(x) := \begin{cases} + 0 & \mathrm{for\ } x < a, \\ + \frac{2(x-a)}{(b-a)(c-a)} & \mathrm{for\ } a \le x < c, \\ + \frac{2}{b-a} & \mathrm{for\ } x = c, \\ + \frac{2(b-x)}{(b-a)(b-c)} & \mathrm{for\ } c < x \le b, \\ + 0 & \mathrm{for\ } b < x. + \end{cases} + + Parameters + ========== + + a : Real number, :math:`a \in \left(-\infty, \infty\right)` + b : Real number, :math:`a < b` + c : Real number, :math:`a \leq c \leq b` + + Returns + ======= + + RandomSymbol + + Examples + ======== + + >>> from sympy.stats import Triangular, density + >>> from sympy import Symbol, pprint + + >>> a = Symbol("a") + >>> b = Symbol("b") + >>> c = Symbol("c") + >>> z = Symbol("z") + + >>> X = Triangular("x", a,b,c) + + >>> pprint(density(X)(z), use_unicode=False) + / -2*a + 2*z + |----------------- for And(a <= z, c > z) + |(-a + b)*(-a + c) + | + | 2 + | ------ for c = z + < -a + b + | + | 2*b - 2*z + |---------------- for And(b >= z, c < z) + |(-a + b)*(b - c) + | + \ 0 otherwise + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Triangular_distribution + .. [2] https://mathworld.wolfram.com/TriangularDistribution.html + + """ + + return rv(name, TriangularDistribution, (a, b, c)) + +#------------------------------------------------------------------------------- +# Uniform distribution --------------------------------------------------------- + + +class UniformDistribution(SingleContinuousDistribution): + _argnames = ('left', 'right') + + @property + def set(self): + return Interval(self.left, self.right) + + @staticmethod + def check(left, right): + _value_check(left < right, "Lower limit should be less than Upper limit.") + + def pdf(self, x): + left, right = self.left, self.right + return Piecewise( + (S.One/(right - left), And(left <= x, x <= right)), + (S.Zero, True) + ) + + def _cdf(self, x): + left, right = self.left, self.right + return Piecewise( + (S.Zero, x < left), + ((x - left)/(right - left), x <= right), + (S.One, True) + ) + + def _characteristic_function(self, t): + left, right = self.left, self.right + return Piecewise(((exp(I*t*right) - exp(I*t*left)) / (I*t*(right - left)), Ne(t, 0)), + (S.One, True)) + + def _moment_generating_function(self, t): + left, right = self.left, self.right + return Piecewise(((exp(t*right) - exp(t*left)) / (t * (right - left)), Ne(t, 0)), + (S.One, True)) + + def expectation(self, expr, var, **kwargs): + kwargs['evaluate'] = True + result = SingleContinuousDistribution.expectation(self, expr, var, **kwargs) + result = result.subs({Max(self.left, self.right): self.right, + Min(self.left, self.right): self.left}) + return result + + +def Uniform(name, left, right): + r""" + Create a continuous random variable with a uniform distribution. + + Explanation + =========== + + The density of the uniform distribution is given by + + .. math:: + f(x) := \begin{cases} + \frac{1}{b - a} & \text{for } x \in [a,b] \\ + 0 & \text{otherwise} + \end{cases} + + with :math:`x \in [a,b]`. + + Parameters + ========== + + a : Real number, :math:`-\infty < a`, the left boundary + b : Real number, :math:`a < b < \infty`, the right boundary + + Returns + ======= + + RandomSymbol + + Examples + ======== + + >>> from sympy.stats import Uniform, density, cdf, E, variance + >>> from sympy import Symbol, simplify + + >>> a = Symbol("a", negative=True) + >>> b = Symbol("b", positive=True) + >>> z = Symbol("z") + + >>> X = Uniform("x", a, b) + + >>> density(X)(z) + Piecewise((1/(-a + b), (b >= z) & (a <= z)), (0, True)) + + >>> cdf(X)(z) + Piecewise((0, a > z), ((-a + z)/(-a + b), b >= z), (1, True)) + + >>> E(X) + a/2 + b/2 + + >>> simplify(variance(X)) + a**2/12 - a*b/6 + b**2/12 + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Uniform_distribution_%28continuous%29 + .. [2] https://mathworld.wolfram.com/UniformDistribution.html + + """ + + return rv(name, UniformDistribution, (left, right)) + +#------------------------------------------------------------------------------- +# UniformSum distribution ------------------------------------------------------ + + +class UniformSumDistribution(SingleContinuousDistribution): + _argnames = ('n',) + + @property + def set(self): + return Interval(0, self.n) + + @staticmethod + def check(n): + _value_check((n > 0, n.is_integer), + "Parameter n must be positive integer.") + + def pdf(self, x): + n = self.n + k = Dummy("k") + return 1/factorial( + n - 1)*Sum((-1)**k*binomial(n, k)*(x - k)**(n - 1), (k, 0, floor(x))) + + def _cdf(self, x): + n = self.n + k = Dummy("k") + return Piecewise((S.Zero, x < 0), + (1/factorial(n)*Sum((-1)**k*binomial(n, k)*(x - k)**(n), + (k, 0, floor(x))), x <= n), + (S.One, True)) + + def _characteristic_function(self, t): + return ((exp(I*t) - 1) / (I*t))**self.n + + def _moment_generating_function(self, t): + return ((exp(t) - 1) / t)**self.n + +def UniformSum(name, n): + r""" + Create a continuous random variable with an Irwin-Hall distribution. + + Explanation + =========== + + The probability distribution function depends on a single parameter + $n$ which is an integer. + + The density of the Irwin-Hall distribution is given by + + .. math :: + f(x) := \frac{1}{(n-1)!}\sum_{k=0}^{\left\lfloor x\right\rfloor}(-1)^k + \binom{n}{k}(x-k)^{n-1} + + Parameters + ========== + + n : A positive integer, `n > 0` + + Returns + ======= + + RandomSymbol + + Examples + ======== + + >>> from sympy.stats import UniformSum, density, cdf + >>> from sympy import Symbol, pprint + + >>> n = Symbol("n", integer=True) + >>> z = Symbol("z") + + >>> X = UniformSum("x", n) + + >>> D = density(X)(z) + >>> pprint(D, use_unicode=False) + floor(z) + ___ + \ ` + \ k n - 1 /n\ + ) (-1) *(-k + z) *| | + / \k/ + /__, + k = 0 + -------------------------------- + (n - 1)! + + >>> cdf(X)(z) + Piecewise((0, z < 0), (Sum((-1)**_k*(-_k + z)**n*binomial(n, _k), + (_k, 0, floor(z)))/factorial(n), n >= z), (1, True)) + + + Compute cdf with specific 'x' and 'n' values as follows : + >>> cdf(UniformSum("x", 5), evaluate=False)(2).doit() + 9/40 + + The argument evaluate=False prevents an attempt at evaluation + of the sum for general n, before the argument 2 is passed. + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Uniform_sum_distribution + .. [2] https://mathworld.wolfram.com/UniformSumDistribution.html + + """ + + return rv(name, UniformSumDistribution, (n, )) + +#------------------------------------------------------------------------------- +# VonMises distribution -------------------------------------------------------- + + +class VonMisesDistribution(SingleContinuousDistribution): + _argnames = ('mu', 'k') + + set = Interval(0, 2*pi) + + @staticmethod + def check(mu, k): + _value_check(k > 0, "k must be positive") + + def pdf(self, x): + mu, k = self.mu, self.k + return exp(k*cos(x-mu)) / (2*pi*besseli(0, k)) + +def VonMises(name, mu, k): + r""" + Create a Continuous Random Variable with a von Mises distribution. + + Explanation + =========== + + The density of the von Mises distribution is given by + + .. math:: + f(x) := \frac{e^{\kappa\cos(x-\mu)}}{2\pi I_0(\kappa)} + + with :math:`x \in [0,2\pi]`. + + Parameters + ========== + + mu : Real number + Measure of location. + k : Real number + Measure of concentration. + + Returns + ======= + + RandomSymbol + + Examples + ======== + + >>> from sympy.stats import VonMises, density + >>> from sympy import Symbol, pprint + + >>> mu = Symbol("mu") + >>> k = Symbol("k", positive=True) + >>> z = Symbol("z") + + >>> X = VonMises("x", mu, k) + + >>> D = density(X)(z) + >>> pprint(D, use_unicode=False) + k*cos(mu - z) + e + ------------------ + 2*pi*besseli(0, k) + + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Von_Mises_distribution + .. [2] https://mathworld.wolfram.com/vonMisesDistribution.html + + """ + + return rv(name, VonMisesDistribution, (mu, k)) + +#------------------------------------------------------------------------------- +# Weibull distribution --------------------------------------------------------- + + +class WeibullDistribution(SingleContinuousDistribution): + _argnames = ('alpha', 'beta') + + set = Interval(0, oo) + + @staticmethod + def check(alpha, beta): + _value_check(alpha > 0, "Alpha must be positive") + _value_check(beta > 0, "Beta must be positive") + + def pdf(self, x): + alpha, beta = self.alpha, self.beta + return beta * (x/alpha)**(beta - 1) * exp(-(x/alpha)**beta) / alpha + + +def Weibull(name, alpha, beta): + r""" + Create a continuous random variable with a Weibull distribution. + + Explanation + =========== + + The density of the Weibull distribution is given by + + .. math:: + f(x) := \begin{cases} + \frac{k}{\lambda}\left(\frac{x}{\lambda}\right)^{k-1} + e^{-(x/\lambda)^{k}} & x\geq0\\ + 0 & x<0 + \end{cases} + + Parameters + ========== + + lambda : Real number, $\lambda > 0$, a scale + k : Real number, $k > 0$, a shape + + Returns + ======= + + RandomSymbol + + Examples + ======== + + >>> from sympy.stats import Weibull, density, E, variance + >>> from sympy import Symbol, simplify + + >>> l = Symbol("lambda", positive=True) + >>> k = Symbol("k", positive=True) + >>> z = Symbol("z") + + >>> X = Weibull("x", l, k) + + >>> density(X)(z) + k*(z/lambda)**(k - 1)*exp(-(z/lambda)**k)/lambda + + >>> simplify(E(X)) + lambda*gamma(1 + 1/k) + + >>> simplify(variance(X)) + lambda**2*(-gamma(1 + 1/k)**2 + gamma(1 + 2/k)) + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Weibull_distribution + .. [2] https://mathworld.wolfram.com/WeibullDistribution.html + + """ + + return rv(name, WeibullDistribution, (alpha, beta)) + +#------------------------------------------------------------------------------- +# Wigner semicircle distribution ----------------------------------------------- + + +class WignerSemicircleDistribution(SingleContinuousDistribution): + _argnames = ('R',) + + @property + def set(self): + return Interval(-self.R, self.R) + + @staticmethod + def check(R): + _value_check(R > 0, "Radius R must be positive.") + + def pdf(self, x): + R = self.R + return 2/(pi*R**2)*sqrt(R**2 - x**2) + + def _characteristic_function(self, t): + return Piecewise((2 * besselj(1, self.R*t) / (self.R*t), Ne(t, 0)), + (S.One, True)) + + def _moment_generating_function(self, t): + return Piecewise((2 * besseli(1, self.R*t) / (self.R*t), Ne(t, 0)), + (S.One, True)) + +def WignerSemicircle(name, R): + r""" + Create a continuous random variable with a Wigner semicircle distribution. + + Explanation + =========== + + The density of the Wigner semicircle distribution is given by + + .. math:: + f(x) := \frac2{\pi R^2}\,\sqrt{R^2-x^2} + + with :math:`x \in [-R,R]`. + + Parameters + ========== + + R : Real number, `R > 0`, the radius + + Returns + ======= + + A RandomSymbol. + + Examples + ======== + + >>> from sympy.stats import WignerSemicircle, density, E + >>> from sympy import Symbol + + >>> R = Symbol("R", positive=True) + >>> z = Symbol("z") + + >>> X = WignerSemicircle("x", R) + + >>> density(X)(z) + 2*sqrt(R**2 - z**2)/(pi*R**2) + + >>> E(X) + 0 + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Wigner_semicircle_distribution + .. [2] https://mathworld.wolfram.com/WignersSemicircleLaw.html + + """ + + return rv(name, WignerSemicircleDistribution, (R,)) diff --git a/lib/python3.12/site-packages/sympy/stats/drv.py b/lib/python3.12/site-packages/sympy/stats/drv.py new file mode 100644 index 0000000000000000000000000000000000000000..13517e0f6dd3cba0362ea5daa148e3d03f23dfd5 --- /dev/null +++ b/lib/python3.12/site-packages/sympy/stats/drv.py @@ -0,0 +1,350 @@ +from sympy.concrete.summations import (Sum, summation) +from sympy.core.basic import Basic +from sympy.core.cache import cacheit +from sympy.core.function import Lambda +from sympy.core.numbers import I +from sympy.core.relational import (Eq, Ne) +from sympy.core.singleton import S +from sympy.core.symbol import (Dummy, symbols) +from sympy.core.sympify import sympify +from sympy.functions.combinatorial.factorials import factorial +from sympy.functions.elementary.exponential import exp +from sympy.functions.elementary.integers import floor +from sympy.functions.elementary.piecewise import Piecewise +from sympy.logic.boolalg import And +from sympy.polys.polytools import poly +from sympy.series.series import series + +from sympy.polys.polyerrors import PolynomialError +from sympy.stats.crv import reduce_rational_inequalities_wrap +from sympy.stats.rv import (NamedArgsMixin, SinglePSpace, SingleDomain, + random_symbols, PSpace, ConditionalDomain, RandomDomain, + ProductDomain, Distribution) +from sympy.stats.symbolic_probability import Probability +from sympy.sets.fancysets import Range, FiniteSet +from sympy.sets.sets import Union +from sympy.sets.contains import Contains +from sympy.utilities import filldedent +from sympy.core.sympify import _sympify + + +class DiscreteDistribution(Distribution): + def __call__(self, *args): + return self.pdf(*args) + + +class SingleDiscreteDistribution(DiscreteDistribution, NamedArgsMixin): + """ Discrete distribution of a single variable. + + Serves as superclass for PoissonDistribution etc.... + + Provides methods for pdf, cdf, and sampling + + See Also: + sympy.stats.crv_types.* + """ + + set = S.Integers + + def __new__(cls, *args): + args = list(map(sympify, args)) + return Basic.__new__(cls, *args) + + @staticmethod + def check(*args): + pass + + @cacheit + def compute_cdf(self, **kwargs): + """ Compute the CDF from the PDF. + + Returns a Lambda. + """ + x = symbols('x', integer=True, cls=Dummy) + z = symbols('z', real=True, cls=Dummy) + left_bound = self.set.inf + + # CDF is integral of PDF from left bound to z + pdf = self.pdf(x) + cdf = summation(pdf, (x, left_bound, floor(z)), **kwargs) + # CDF Ensure that CDF left of left_bound is zero + cdf = Piecewise((cdf, z >= left_bound), (0, True)) + return Lambda(z, cdf) + + def _cdf(self, x): + return None + + def cdf(self, x, **kwargs): + """ Cumulative density function """ + if not kwargs: + cdf = self._cdf(x) + if cdf is not None: + return cdf + return self.compute_cdf(**kwargs)(x) + + @cacheit + def compute_characteristic_function(self, **kwargs): + """ Compute the characteristic function from the PDF. + + Returns a Lambda. + """ + x, t = symbols('x, t', real=True, cls=Dummy) + pdf = self.pdf(x) + cf = summation(exp(I*t*x)*pdf, (x, self.set.inf, self.set.sup)) + return Lambda(t, cf) + + def _characteristic_function(self, t): + return None + + def characteristic_function(self, t, **kwargs): + """ Characteristic function """ + if not kwargs: + cf = self._characteristic_function(t) + if cf is not None: + return cf + return self.compute_characteristic_function(**kwargs)(t) + + @cacheit + def compute_moment_generating_function(self, **kwargs): + t = Dummy('t', real=True) + x = Dummy('x', integer=True) + pdf = self.pdf(x) + mgf = summation(exp(t*x)*pdf, (x, self.set.inf, self.set.sup)) + return Lambda(t, mgf) + + def _moment_generating_function(self, t): + return None + + def moment_generating_function(self, t, **kwargs): + if not kwargs: + mgf = self._moment_generating_function(t) + if mgf is not None: + return mgf + return self.compute_moment_generating_function(**kwargs)(t) + + @cacheit + def compute_quantile(self, **kwargs): + """ Compute the Quantile from the PDF. + + Returns a Lambda. + """ + x = Dummy('x', integer=True) + p = Dummy('p', real=True) + left_bound = self.set.inf + pdf = self.pdf(x) + cdf = summation(pdf, (x, left_bound, x), **kwargs) + set = ((x, p <= cdf), ) + return Lambda(p, Piecewise(*set)) + + def _quantile(self, x): + return None + + def quantile(self, x, **kwargs): + """ Cumulative density function """ + if not kwargs: + quantile = self._quantile(x) + if quantile is not None: + return quantile + return self.compute_quantile(**kwargs)(x) + + def expectation(self, expr, var, evaluate=True, **kwargs): + """ Expectation of expression over distribution """ + # TODO: support discrete sets with non integer stepsizes + + if evaluate: + try: + p = poly(expr, var) + + t = Dummy('t', real=True) + + mgf = self.moment_generating_function(t) + deg = p.degree() + taylor = poly(series(mgf, t, 0, deg + 1).removeO(), t) + result = 0 + for k in range(deg+1): + result += p.coeff_monomial(var ** k) * taylor.coeff_monomial(t ** k) * factorial(k) + + return result + + except PolynomialError: + return summation(expr * self.pdf(var), + (var, self.set.inf, self.set.sup), **kwargs) + + else: + return Sum(expr * self.pdf(var), + (var, self.set.inf, self.set.sup), **kwargs) + + def __call__(self, *args): + return self.pdf(*args) + + +class DiscreteDomain(RandomDomain): + """ + A domain with discrete support with step size one. + Represented using symbols and Range. + """ + is_Discrete = True + +class SingleDiscreteDomain(DiscreteDomain, SingleDomain): + def as_boolean(self): + return Contains(self.symbol, self.set) + + +class ConditionalDiscreteDomain(DiscreteDomain, ConditionalDomain): + """ + Domain with discrete support of step size one, that is restricted by + some condition. + """ + @property + def set(self): + rv = self.symbols + if len(self.symbols) > 1: + raise NotImplementedError(filldedent(''' + Multivariate conditional domains are not yet implemented.''')) + rv = list(rv)[0] + return reduce_rational_inequalities_wrap(self.condition, + rv).intersect(self.fulldomain.set) + + +class DiscretePSpace(PSpace): + is_real = True + is_Discrete = True + + @property + def pdf(self): + return self.density(*self.symbols) + + def where(self, condition): + rvs = random_symbols(condition) + assert all(r.symbol in self.symbols for r in rvs) + if len(rvs) > 1: + raise NotImplementedError(filldedent('''Multivariate discrete + random variables are not yet supported.''')) + conditional_domain = reduce_rational_inequalities_wrap(condition, + rvs[0]) + conditional_domain = conditional_domain.intersect(self.domain.set) + return SingleDiscreteDomain(rvs[0].symbol, conditional_domain) + + def probability(self, condition): + complement = isinstance(condition, Ne) + if complement: + condition = Eq(condition.args[0], condition.args[1]) + try: + _domain = self.where(condition).set + if condition == False or _domain is S.EmptySet: + return S.Zero + if condition == True or _domain == self.domain.set: + return S.One + prob = self.eval_prob(_domain) + except NotImplementedError: + from sympy.stats.rv import density + expr = condition.lhs - condition.rhs + dens = density(expr) + if not isinstance(dens, DiscreteDistribution): + from sympy.stats.drv_types import DiscreteDistributionHandmade + dens = DiscreteDistributionHandmade(dens) + z = Dummy('z', real=True) + space = SingleDiscretePSpace(z, dens) + prob = space.probability(condition.__class__(space.value, 0)) + if prob is None: + prob = Probability(condition) + return prob if not complement else S.One - prob + + def eval_prob(self, _domain): + sym = list(self.symbols)[0] + if isinstance(_domain, Range): + n = symbols('n', integer=True) + inf, sup, step = (r for r in _domain.args) + summand = ((self.pdf).replace( + sym, n*step)) + rv = summation(summand, + (n, inf/step, (sup)/step - 1)).doit() + return rv + elif isinstance(_domain, FiniteSet): + pdf = Lambda(sym, self.pdf) + rv = sum(pdf(x) for x in _domain) + return rv + elif isinstance(_domain, Union): + rv = sum(self.eval_prob(x) for x in _domain.args) + return rv + + def conditional_space(self, condition): + # XXX: Converting from set to tuple. The order matters to Lambda + # though so we should be starting with a set... + density = Lambda(tuple(self.symbols), self.pdf/self.probability(condition)) + condition = condition.xreplace({rv: rv.symbol for rv in self.values}) + domain = ConditionalDiscreteDomain(self.domain, condition) + return DiscretePSpace(domain, density) + +class ProductDiscreteDomain(ProductDomain, DiscreteDomain): + def as_boolean(self): + return And(*[domain.as_boolean for domain in self.domains]) + +class SingleDiscretePSpace(DiscretePSpace, SinglePSpace): + """ Discrete probability space over a single univariate variable """ + is_real = True + + @property + def set(self): + return self.distribution.set + + @property + def domain(self): + return SingleDiscreteDomain(self.symbol, self.set) + + def sample(self, size=(), library='scipy', seed=None): + """ + Internal sample method. + + Returns dictionary mapping RandomSymbol to realization value. + """ + return {self.value: self.distribution.sample(size, library=library, seed=seed)} + + def compute_expectation(self, expr, rvs=None, evaluate=True, **kwargs): + rvs = rvs or (self.value,) + if self.value not in rvs: + return expr + + expr = _sympify(expr) + expr = expr.xreplace({rv: rv.symbol for rv in rvs}) + + x = self.value.symbol + try: + return self.distribution.expectation(expr, x, evaluate=evaluate, + **kwargs) + except NotImplementedError: + return Sum(expr * self.pdf, (x, self.set.inf, self.set.sup), + **kwargs) + + def compute_cdf(self, expr, **kwargs): + if expr == self.value: + x = Dummy("x", real=True) + return Lambda(x, self.distribution.cdf(x, **kwargs)) + else: + raise NotImplementedError() + + def compute_density(self, expr, **kwargs): + if expr == self.value: + return self.distribution + raise NotImplementedError() + + def compute_characteristic_function(self, expr, **kwargs): + if expr == self.value: + t = Dummy("t", real=True) + return Lambda(t, self.distribution.characteristic_function(t, **kwargs)) + else: + raise NotImplementedError() + + def compute_moment_generating_function(self, expr, **kwargs): + if expr == self.value: + t = Dummy("t", real=True) + return Lambda(t, self.distribution.moment_generating_function(t, **kwargs)) + else: + raise NotImplementedError() + + def compute_quantile(self, expr, **kwargs): + if expr == self.value: + p = Dummy("p", real=True) + return Lambda(p, self.distribution.quantile(p, **kwargs)) + else: + raise NotImplementedError() diff --git a/lib/python3.12/site-packages/sympy/stats/drv_types.py b/lib/python3.12/site-packages/sympy/stats/drv_types.py new file mode 100644 index 0000000000000000000000000000000000000000..a2ce892168bdbbe24d8f8a5586e7295e07df25ea --- /dev/null +++ b/lib/python3.12/site-packages/sympy/stats/drv_types.py @@ -0,0 +1,835 @@ +""" + +Contains +======== +FlorySchulz +Geometric +Hermite +Logarithmic +NegativeBinomial +Poisson +Skellam +YuleSimon +Zeta +""" + + + +from sympy.concrete.summations import Sum +from sympy.core.basic import Basic +from sympy.core.function import Lambda +from sympy.core.numbers import I +from sympy.core.relational import Eq +from sympy.core.singleton import S +from sympy.core.symbol import Dummy +from sympy.core.sympify import sympify +from sympy.functions.combinatorial.factorials import (binomial, factorial) +from sympy.functions.elementary.exponential import (exp, log) +from sympy.functions.elementary.integers import floor +from sympy.functions.elementary.miscellaneous import sqrt +from sympy.functions.elementary.piecewise import Piecewise +from sympy.functions.special.bessel import besseli +from sympy.functions.special.beta_functions import beta +from sympy.functions.special.hyper import hyper +from sympy.functions.special.zeta_functions import (polylog, zeta) +from sympy.stats.drv import SingleDiscreteDistribution, SingleDiscretePSpace +from sympy.stats.rv import _value_check, is_random + + +__all__ = ['FlorySchulz', +'Geometric', +'Hermite', +'Logarithmic', +'NegativeBinomial', +'Poisson', +'Skellam', +'YuleSimon', +'Zeta' +] + + +def rv(symbol, cls, *args, **kwargs): + args = list(map(sympify, args)) + dist = cls(*args) + if kwargs.pop('check', True): + dist.check(*args) + pspace = SingleDiscretePSpace(symbol, dist) + if any(is_random(arg) for arg in args): + from sympy.stats.compound_rv import CompoundPSpace, CompoundDistribution + pspace = CompoundPSpace(symbol, CompoundDistribution(dist)) + return pspace.value + + +class DiscreteDistributionHandmade(SingleDiscreteDistribution): + _argnames = ('pdf',) + + def __new__(cls, pdf, set=S.Integers): + return Basic.__new__(cls, pdf, set) + + @property + def set(self): + return self.args[1] + + @staticmethod + def check(pdf, set): + x = Dummy('x') + val = Sum(pdf(x), (x, set._inf, set._sup)).doit() + _value_check(Eq(val, 1) != S.false, "The pdf is incorrect on the given set.") + + + +def DiscreteRV(symbol, density, set=S.Integers, **kwargs): + """ + Create a Discrete Random Variable given the following: + + Parameters + ========== + + symbol : Symbol + Represents name of the random variable. + density : Expression containing symbol + Represents probability density function. + set : set + Represents the region where the pdf is valid, by default is real line. + check : bool + If True, it will check whether the given density + integrates to 1 over the given set. If False, it + will not perform this check. Default is False. + + Examples + ======== + + >>> from sympy.stats import DiscreteRV, P, E + >>> from sympy import Rational, Symbol + >>> x = Symbol('x') + >>> n = 10 + >>> density = Rational(1, 10) + >>> X = DiscreteRV(x, density, set=set(range(n))) + >>> E(X) + 9/2 + >>> P(X>3) + 3/5 + + Returns + ======= + + RandomSymbol + + """ + set = sympify(set) + pdf = Piecewise((density, set.as_relational(symbol)), (0, True)) + pdf = Lambda(symbol, pdf) + # have a default of False while `rv` should have a default of True + kwargs['check'] = kwargs.pop('check', False) + return rv(symbol.name, DiscreteDistributionHandmade, pdf, set, **kwargs) + + +#------------------------------------------------------------------------------- +# Flory-Schulz distribution ------------------------------------------------------------ + +class FlorySchulzDistribution(SingleDiscreteDistribution): + _argnames = ('a',) + set = S.Naturals + + @staticmethod + def check(a): + _value_check((0 < a, a < 1), "a must be between 0 and 1") + + def pdf(self, k): + a = self.a + return (a**2 * k * (1 - a)**(k - 1)) + + def _characteristic_function(self, t): + a = self.a + return a**2*exp(I*t)/((1 + (a - 1)*exp(I*t))**2) + + def _moment_generating_function(self, t): + a = self.a + return a**2*exp(t)/((1 + (a - 1)*exp(t))**2) + + +def FlorySchulz(name, a): + r""" + Create a discrete random variable with a FlorySchulz distribution. + + The density of the FlorySchulz distribution is given by + + .. math:: + f(k) := (a^2) k (1 - a)^{k-1} + + Parameters + ========== + + a : A real number between 0 and 1 + + Returns + ======= + + RandomSymbol + + Examples + ======== + + >>> from sympy.stats import density, E, variance, FlorySchulz + >>> from sympy import Symbol, S + + >>> a = S.One / 5 + >>> z = Symbol("z") + + >>> X = FlorySchulz("x", a) + + >>> density(X)(z) + (5/4)**(1 - z)*z/25 + + >>> E(X) + 9 + + >>> variance(X) + 40 + + References + ========== + + https://en.wikipedia.org/wiki/Flory%E2%80%93Schulz_distribution + """ + return rv(name, FlorySchulzDistribution, a) + + +#------------------------------------------------------------------------------- +# Geometric distribution ------------------------------------------------------------ + +class GeometricDistribution(SingleDiscreteDistribution): + _argnames = ('p',) + set = S.Naturals + + @staticmethod + def check(p): + _value_check((0 < p, p <= 1), "p must be between 0 and 1") + + def pdf(self, k): + return (1 - self.p)**(k - 1) * self.p + + def _characteristic_function(self, t): + p = self.p + return p * exp(I*t) / (1 - (1 - p)*exp(I*t)) + + def _moment_generating_function(self, t): + p = self.p + return p * exp(t) / (1 - (1 - p) * exp(t)) + + +def Geometric(name, p): + r""" + Create a discrete random variable with a Geometric distribution. + + Explanation + =========== + + The density of the Geometric distribution is given by + + .. math:: + f(k) := p (1 - p)^{k - 1} + + Parameters + ========== + + p : A probability between 0 and 1 + + Returns + ======= + + RandomSymbol + + Examples + ======== + + >>> from sympy.stats import Geometric, density, E, variance + >>> from sympy import Symbol, S + + >>> p = S.One / 5 + >>> z = Symbol("z") + + >>> X = Geometric("x", p) + + >>> density(X)(z) + (5/4)**(1 - z)/5 + + >>> E(X) + 5 + + >>> variance(X) + 20 + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Geometric_distribution + .. [2] https://mathworld.wolfram.com/GeometricDistribution.html + + """ + return rv(name, GeometricDistribution, p) + + +#------------------------------------------------------------------------------- +# Hermite distribution --------------------------------------------------------- + + +class HermiteDistribution(SingleDiscreteDistribution): + _argnames = ('a1', 'a2') + set = S.Naturals0 + + @staticmethod + def check(a1, a2): + _value_check(a1.is_nonnegative, 'Parameter a1 must be >= 0.') + _value_check(a2.is_nonnegative, 'Parameter a2 must be >= 0.') + + def pdf(self, k): + a1, a2 = self.a1, self.a2 + term1 = exp(-(a1 + a2)) + j = Dummy("j", integer=True) + num = a1**(k - 2*j) * a2**j + den = factorial(k - 2*j) * factorial(j) + return term1 * Sum(num/den, (j, 0, k//2)).doit() + + def _moment_generating_function(self, t): + a1, a2 = self.a1, self.a2 + term1 = a1 * (exp(t) - 1) + term2 = a2 * (exp(2*t) - 1) + return exp(term1 + term2) + + def _characteristic_function(self, t): + a1, a2 = self.a1, self.a2 + term1 = a1 * (exp(I*t) - 1) + term2 = a2 * (exp(2*I*t) - 1) + return exp(term1 + term2) + +def Hermite(name, a1, a2): + r""" + Create a discrete random variable with a Hermite distribution. + + Explanation + =========== + + The density of the Hermite distribution is given by + + .. math:: + f(x):= e^{-a_1 -a_2}\sum_{j=0}^{\left \lfloor x/2 \right \rfloor} + \frac{a_{1}^{x-2j}a_{2}^{j}}{(x-2j)!j!} + + Parameters + ========== + + a1 : A Positive number greater than equal to 0. + a2 : A Positive number greater than equal to 0. + + Returns + ======= + + RandomSymbol + + Examples + ======== + + >>> from sympy.stats import Hermite, density, E, variance + >>> from sympy import Symbol + + >>> a1 = Symbol("a1", positive=True) + >>> a2 = Symbol("a2", positive=True) + >>> x = Symbol("x") + + >>> H = Hermite("H", a1=5, a2=4) + + >>> density(H)(2) + 33*exp(-9)/2 + + >>> E(H) + 13 + + >>> variance(H) + 21 + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Hermite_distribution + + """ + + return rv(name, HermiteDistribution, a1, a2) + + +#------------------------------------------------------------------------------- +# Logarithmic distribution ------------------------------------------------------------ + +class LogarithmicDistribution(SingleDiscreteDistribution): + _argnames = ('p',) + + set = S.Naturals + + @staticmethod + def check(p): + _value_check((p > 0, p < 1), "p should be between 0 and 1") + + def pdf(self, k): + p = self.p + return (-1) * p**k / (k * log(1 - p)) + + def _characteristic_function(self, t): + p = self.p + return log(1 - p * exp(I*t)) / log(1 - p) + + def _moment_generating_function(self, t): + p = self.p + return log(1 - p * exp(t)) / log(1 - p) + + +def Logarithmic(name, p): + r""" + Create a discrete random variable with a Logarithmic distribution. + + Explanation + =========== + + The density of the Logarithmic distribution is given by + + .. math:: + f(k) := \frac{-p^k}{k \ln{(1 - p)}} + + Parameters + ========== + + p : A value between 0 and 1 + + Returns + ======= + + RandomSymbol + + Examples + ======== + + >>> from sympy.stats import Logarithmic, density, E, variance + >>> from sympy import Symbol, S + + >>> p = S.One / 5 + >>> z = Symbol("z") + + >>> X = Logarithmic("x", p) + + >>> density(X)(z) + -1/(5**z*z*log(4/5)) + + >>> E(X) + -1/(-4*log(5) + 8*log(2)) + + >>> variance(X) + -1/((-4*log(5) + 8*log(2))*(-2*log(5) + 4*log(2))) + 1/(-64*log(2)*log(5) + 64*log(2)**2 + 16*log(5)**2) - 10/(-32*log(5) + 64*log(2)) + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Logarithmic_distribution + .. [2] https://mathworld.wolfram.com/LogarithmicDistribution.html + + """ + return rv(name, LogarithmicDistribution, p) + + +#------------------------------------------------------------------------------- +# Negative binomial distribution ------------------------------------------------------------ + +class NegativeBinomialDistribution(SingleDiscreteDistribution): + _argnames = ('r', 'p') + set = S.Naturals0 + + @staticmethod + def check(r, p): + _value_check(r > 0, 'r should be positive') + _value_check((p > 0, p < 1), 'p should be between 0 and 1') + + def pdf(self, k): + r = self.r + p = self.p + + return binomial(k + r - 1, k) * (1 - p)**r * p**k + + def _characteristic_function(self, t): + r = self.r + p = self.p + + return ((1 - p) / (1 - p * exp(I*t)))**r + + def _moment_generating_function(self, t): + r = self.r + p = self.p + + return ((1 - p) / (1 - p * exp(t)))**r + +def NegativeBinomial(name, r, p): + r""" + Create a discrete random variable with a Negative Binomial distribution. + + Explanation + =========== + + The density of the Negative Binomial distribution is given by + + .. math:: + f(k) := \binom{k + r - 1}{k} (1 - p)^r p^k + + Parameters + ========== + + r : A positive value + p : A value between 0 and 1 + + Returns + ======= + + RandomSymbol + + Examples + ======== + + >>> from sympy.stats import NegativeBinomial, density, E, variance + >>> from sympy import Symbol, S + + >>> r = 5 + >>> p = S.One / 5 + >>> z = Symbol("z") + + >>> X = NegativeBinomial("x", r, p) + + >>> density(X)(z) + 1024*binomial(z + 4, z)/(3125*5**z) + + >>> E(X) + 5/4 + + >>> variance(X) + 25/16 + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Negative_binomial_distribution + .. [2] https://mathworld.wolfram.com/NegativeBinomialDistribution.html + + """ + return rv(name, NegativeBinomialDistribution, r, p) + + +#------------------------------------------------------------------------------- +# Poisson distribution ------------------------------------------------------------ + +class PoissonDistribution(SingleDiscreteDistribution): + _argnames = ('lamda',) + + set = S.Naturals0 + + @staticmethod + def check(lamda): + _value_check(lamda > 0, "Lambda must be positive") + + def pdf(self, k): + return self.lamda**k / factorial(k) * exp(-self.lamda) + + def _characteristic_function(self, t): + return exp(self.lamda * (exp(I*t) - 1)) + + def _moment_generating_function(self, t): + return exp(self.lamda * (exp(t) - 1)) + + +def Poisson(name, lamda): + r""" + Create a discrete random variable with a Poisson distribution. + + Explanation + =========== + + The density of the Poisson distribution is given by + + .. math:: + f(k) := \frac{\lambda^{k} e^{- \lambda}}{k!} + + Parameters + ========== + + lamda : Positive number, a rate + + Returns + ======= + + RandomSymbol + + Examples + ======== + + >>> from sympy.stats import Poisson, density, E, variance + >>> from sympy import Symbol, simplify + + >>> rate = Symbol("lambda", positive=True) + >>> z = Symbol("z") + + >>> X = Poisson("x", rate) + + >>> density(X)(z) + lambda**z*exp(-lambda)/factorial(z) + + >>> E(X) + lambda + + >>> simplify(variance(X)) + lambda + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Poisson_distribution + .. [2] https://mathworld.wolfram.com/PoissonDistribution.html + + """ + return rv(name, PoissonDistribution, lamda) + + +# ----------------------------------------------------------------------------- +# Skellam distribution -------------------------------------------------------- + + +class SkellamDistribution(SingleDiscreteDistribution): + _argnames = ('mu1', 'mu2') + set = S.Integers + + @staticmethod + def check(mu1, mu2): + _value_check(mu1 >= 0, 'Parameter mu1 must be >= 0') + _value_check(mu2 >= 0, 'Parameter mu2 must be >= 0') + + def pdf(self, k): + (mu1, mu2) = (self.mu1, self.mu2) + term1 = exp(-(mu1 + mu2)) * (mu1 / mu2) ** (k / 2) + term2 = besseli(k, 2 * sqrt(mu1 * mu2)) + return term1 * term2 + + def _cdf(self, x): + raise NotImplementedError( + "Skellam doesn't have closed form for the CDF.") + + def _characteristic_function(self, t): + (mu1, mu2) = (self.mu1, self.mu2) + return exp(-(mu1 + mu2) + mu1 * exp(I * t) + mu2 * exp(-I * t)) + + def _moment_generating_function(self, t): + (mu1, mu2) = (self.mu1, self.mu2) + return exp(-(mu1 + mu2) + mu1 * exp(t) + mu2 * exp(-t)) + + +def Skellam(name, mu1, mu2): + r""" + Create a discrete random variable with a Skellam distribution. + + Explanation + =========== + + The Skellam is the distribution of the difference N1 - N2 + of two statistically independent random variables N1 and N2 + each Poisson-distributed with respective expected values mu1 and mu2. + + The density of the Skellam distribution is given by + + .. math:: + f(k) := e^{-(\mu_1+\mu_2)}(\frac{\mu_1}{\mu_2})^{k/2}I_k(2\sqrt{\mu_1\mu_2}) + + Parameters + ========== + + mu1 : A non-negative value + mu2 : A non-negative value + + Returns + ======= + + RandomSymbol + + Examples + ======== + + >>> from sympy.stats import Skellam, density, E, variance + >>> from sympy import Symbol, pprint + + >>> z = Symbol("z", integer=True) + >>> mu1 = Symbol("mu1", positive=True) + >>> mu2 = Symbol("mu2", positive=True) + >>> X = Skellam("x", mu1, mu2) + + >>> pprint(density(X)(z), use_unicode=False) + z + - + 2 + /mu1\ -mu1 - mu2 / _____ _____\ + |---| *e *besseli\z, 2*\/ mu1 *\/ mu2 / + \mu2/ + >>> E(X) + mu1 - mu2 + >>> variance(X).expand() + mu1 + mu2 + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Skellam_distribution + + """ + return rv(name, SkellamDistribution, mu1, mu2) + + +#------------------------------------------------------------------------------- +# Yule-Simon distribution ------------------------------------------------------------ + +class YuleSimonDistribution(SingleDiscreteDistribution): + _argnames = ('rho',) + set = S.Naturals + + @staticmethod + def check(rho): + _value_check(rho > 0, 'rho should be positive') + + def pdf(self, k): + rho = self.rho + return rho * beta(k, rho + 1) + + def _cdf(self, x): + return Piecewise((1 - floor(x) * beta(floor(x), self.rho + 1), x >= 1), (0, True)) + + def _characteristic_function(self, t): + rho = self.rho + return rho * hyper((1, 1), (rho + 2,), exp(I*t)) * exp(I*t) / (rho + 1) + + def _moment_generating_function(self, t): + rho = self.rho + return rho * hyper((1, 1), (rho + 2,), exp(t)) * exp(t) / (rho + 1) + + +def YuleSimon(name, rho): + r""" + Create a discrete random variable with a Yule-Simon distribution. + + Explanation + =========== + + The density of the Yule-Simon distribution is given by + + .. math:: + f(k) := \rho B(k, \rho + 1) + + Parameters + ========== + + rho : A positive value + + Returns + ======= + + RandomSymbol + + Examples + ======== + + >>> from sympy.stats import YuleSimon, density, E, variance + >>> from sympy import Symbol, simplify + + >>> p = 5 + >>> z = Symbol("z") + + >>> X = YuleSimon("x", p) + + >>> density(X)(z) + 5*beta(z, 6) + + >>> simplify(E(X)) + 5/4 + + >>> simplify(variance(X)) + 25/48 + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Yule%E2%80%93Simon_distribution + + """ + return rv(name, YuleSimonDistribution, rho) + + +#------------------------------------------------------------------------------- +# Zeta distribution ------------------------------------------------------------ + +class ZetaDistribution(SingleDiscreteDistribution): + _argnames = ('s',) + set = S.Naturals + + @staticmethod + def check(s): + _value_check(s > 1, 's should be greater than 1') + + def pdf(self, k): + s = self.s + return 1 / (k**s * zeta(s)) + + def _characteristic_function(self, t): + return polylog(self.s, exp(I*t)) / zeta(self.s) + + def _moment_generating_function(self, t): + return polylog(self.s, exp(t)) / zeta(self.s) + + +def Zeta(name, s): + r""" + Create a discrete random variable with a Zeta distribution. + + Explanation + =========== + + The density of the Zeta distribution is given by + + .. math:: + f(k) := \frac{1}{k^s \zeta{(s)}} + + Parameters + ========== + + s : A value greater than 1 + + Returns + ======= + + RandomSymbol + + Examples + ======== + + >>> from sympy.stats import Zeta, density, E, variance + >>> from sympy import Symbol + + >>> s = 5 + >>> z = Symbol("z") + + >>> X = Zeta("x", s) + + >>> density(X)(z) + 1/(z**5*zeta(5)) + + >>> E(X) + pi**4/(90*zeta(5)) + + >>> variance(X) + -pi**8/(8100*zeta(5)**2) + zeta(3)/zeta(5) + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Zeta_distribution + + """ + return rv(name, ZetaDistribution, s) diff --git a/lib/python3.12/site-packages/sympy/stats/error_prop.py b/lib/python3.12/site-packages/sympy/stats/error_prop.py new file mode 100644 index 0000000000000000000000000000000000000000..e6cacb894307fe60cbf096c7760e6ed57f385a91 --- /dev/null +++ b/lib/python3.12/site-packages/sympy/stats/error_prop.py @@ -0,0 +1,100 @@ +"""Tools for arithmetic error propagation.""" + +from itertools import repeat, combinations + +from sympy.core.add import Add +from sympy.core.mul import Mul +from sympy.core.power import Pow +from sympy.core.singleton import S +from sympy.core.symbol import Symbol +from sympy.functions.elementary.exponential import exp +from sympy.simplify.simplify import simplify +from sympy.stats.symbolic_probability import RandomSymbol, Variance, Covariance +from sympy.stats.rv import is_random + +_arg0_or_var = lambda var: var.args[0] if len(var.args) > 0 else var + + +def variance_prop(expr, consts=(), include_covar=False): + r"""Symbolically propagates variance (`\sigma^2`) for expressions. + This is computed as as seen in [1]_. + + Parameters + ========== + + expr : Expr + A SymPy expression to compute the variance for. + consts : sequence of Symbols, optional + Represents symbols that are known constants in the expr, + and thus have zero variance. All symbols not in consts are + assumed to be variant. + include_covar : bool, optional + Flag for whether or not to include covariances, default=False. + + Returns + ======= + + var_expr : Expr + An expression for the total variance of the expr. + The variance for the original symbols (e.g. x) are represented + via instance of the Variance symbol (e.g. Variance(x)). + + Examples + ======== + + >>> from sympy import symbols, exp + >>> from sympy.stats.error_prop import variance_prop + >>> x, y = symbols('x y') + + >>> variance_prop(x + y) + Variance(x) + Variance(y) + + >>> variance_prop(x * y) + x**2*Variance(y) + y**2*Variance(x) + + >>> variance_prop(exp(2*x)) + 4*exp(4*x)*Variance(x) + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Propagation_of_uncertainty + + """ + args = expr.args + if len(args) == 0: + if expr in consts: + return S.Zero + elif is_random(expr): + return Variance(expr).doit() + elif isinstance(expr, Symbol): + return Variance(RandomSymbol(expr)).doit() + else: + return S.Zero + nargs = len(args) + var_args = list(map(variance_prop, args, repeat(consts, nargs), + repeat(include_covar, nargs))) + if isinstance(expr, Add): + var_expr = Add(*var_args) + if include_covar: + terms = [2 * Covariance(_arg0_or_var(x), _arg0_or_var(y)).expand() \ + for x, y in combinations(var_args, 2)] + var_expr += Add(*terms) + elif isinstance(expr, Mul): + terms = [v/a**2 for a, v in zip(args, var_args)] + var_expr = simplify(expr**2 * Add(*terms)) + if include_covar: + terms = [2*Covariance(_arg0_or_var(x), _arg0_or_var(y)).expand()/(a*b) \ + for (a, b), (x, y) in zip(combinations(args, 2), + combinations(var_args, 2))] + var_expr += Add(*terms) + elif isinstance(expr, Pow): + b = args[1] + v = var_args[0] * (expr * b / args[0])**2 + var_expr = simplify(v) + elif isinstance(expr, exp): + var_expr = simplify(var_args[0] * expr**2) + else: + # unknown how to proceed, return variance of whole expr. + var_expr = Variance(expr) + return var_expr diff --git a/lib/python3.12/site-packages/sympy/stats/frv.py b/lib/python3.12/site-packages/sympy/stats/frv.py new file mode 100644 index 0000000000000000000000000000000000000000..498d7e4006b2b8db306a0905ed67578021e220a8 --- /dev/null +++ b/lib/python3.12/site-packages/sympy/stats/frv.py @@ -0,0 +1,512 @@ +""" +Finite Discrete Random Variables Module + +See Also +======== +sympy.stats.frv_types +sympy.stats.rv +sympy.stats.crv +""" +from itertools import product + +from sympy.concrete.summations import Sum +from sympy.core.basic import Basic +from sympy.core.cache import cacheit +from sympy.core.function import Lambda +from sympy.core.mul import Mul +from sympy.core.numbers import (I, nan) +from sympy.core.relational import Eq +from sympy.core.singleton import S +from sympy.core.symbol import (Dummy, Symbol) +from sympy.core.sympify import sympify +from sympy.functions.elementary.exponential import exp +from sympy.functions.elementary.piecewise import Piecewise +from sympy.logic.boolalg import (And, Or) +from sympy.sets.sets import Intersection +from sympy.core.containers import Dict +from sympy.core.logic import Logic +from sympy.core.relational import Relational +from sympy.core.sympify import _sympify +from sympy.sets.sets import FiniteSet +from sympy.stats.rv import (RandomDomain, ProductDomain, ConditionalDomain, + PSpace, IndependentProductPSpace, SinglePSpace, random_symbols, + sumsets, rv_subs, NamedArgsMixin, Density, Distribution) + + +class FiniteDensity(dict): + """ + A domain with Finite Density. + """ + def __call__(self, item): + """ + Make instance of a class callable. + + If item belongs to current instance of a class, return it. + + Otherwise, return 0. + """ + item = sympify(item) + if item in self: + return self[item] + else: + return 0 + + @property + def dict(self): + """ + Return item as dictionary. + """ + return dict(self) + +class FiniteDomain(RandomDomain): + """ + A domain with discrete finite support + + Represented using a FiniteSet. + """ + is_Finite = True + + @property + def symbols(self): + return FiniteSet(sym for sym, val in self.elements) + + @property + def elements(self): + return self.args[0] + + @property + def dict(self): + return FiniteSet(*[Dict(dict(el)) for el in self.elements]) + + def __contains__(self, other): + return other in self.elements + + def __iter__(self): + return self.elements.__iter__() + + def as_boolean(self): + return Or(*[And(*[Eq(sym, val) for sym, val in item]) for item in self]) + + +class SingleFiniteDomain(FiniteDomain): + """ + A FiniteDomain over a single symbol/set + + Example: The possibilities of a *single* die roll. + """ + + def __new__(cls, symbol, set): + if not isinstance(set, FiniteSet) and \ + not isinstance(set, Intersection): + set = FiniteSet(*set) + return Basic.__new__(cls, symbol, set) + + @property + def symbol(self): + return self.args[0] + + @property + def symbols(self): + return FiniteSet(self.symbol) + + @property + def set(self): + return self.args[1] + + @property + def elements(self): + return FiniteSet(*[frozenset(((self.symbol, elem), )) for elem in self.set]) + + def __iter__(self): + return (frozenset(((self.symbol, elem),)) for elem in self.set) + + def __contains__(self, other): + sym, val = tuple(other)[0] + return sym == self.symbol and val in self.set + + +class ProductFiniteDomain(ProductDomain, FiniteDomain): + """ + A Finite domain consisting of several other FiniteDomains + + Example: The possibilities of the rolls of three independent dice + """ + + def __iter__(self): + proditer = product(*self.domains) + return (sumsets(items) for items in proditer) + + @property + def elements(self): + return FiniteSet(*self) + + +class ConditionalFiniteDomain(ConditionalDomain, ProductFiniteDomain): + """ + A FiniteDomain that has been restricted by a condition + + Example: The possibilities of a die roll under the condition that the + roll is even. + """ + + def __new__(cls, domain, condition): + """ + Create a new instance of ConditionalFiniteDomain class + """ + if condition is True: + return domain + cond = rv_subs(condition) + return Basic.__new__(cls, domain, cond) + + def _test(self, elem): + """ + Test the value. If value is boolean, return it. If value is equality + relational (two objects are equal), return it with left-hand side + being equal to right-hand side. Otherwise, raise ValueError exception. + """ + val = self.condition.xreplace(dict(elem)) + if val in [True, False]: + return val + elif val.is_Equality: + return val.lhs == val.rhs + raise ValueError("Undecidable if %s" % str(val)) + + def __contains__(self, other): + return other in self.fulldomain and self._test(other) + + def __iter__(self): + return (elem for elem in self.fulldomain if self._test(elem)) + + @property + def set(self): + if isinstance(self.fulldomain, SingleFiniteDomain): + return FiniteSet(*[elem for elem in self.fulldomain.set + if frozenset(((self.fulldomain.symbol, elem),)) in self]) + else: + raise NotImplementedError( + "Not implemented on multi-dimensional conditional domain") + + def as_boolean(self): + return FiniteDomain.as_boolean(self) + + +class SingleFiniteDistribution(Distribution, NamedArgsMixin): + def __new__(cls, *args): + args = list(map(sympify, args)) + return Basic.__new__(cls, *args) + + @staticmethod + def check(*args): + pass + + @property # type: ignore + @cacheit + def dict(self): + if self.is_symbolic: + return Density(self) + return {k: self.pmf(k) for k in self.set} + + def pmf(self, *args): # to be overridden by specific distribution + raise NotImplementedError() + + @property + def set(self): # to be overridden by specific distribution + raise NotImplementedError() + + values = property(lambda self: self.dict.values) + items = property(lambda self: self.dict.items) + is_symbolic = property(lambda self: False) + __iter__ = property(lambda self: self.dict.__iter__) + __getitem__ = property(lambda self: self.dict.__getitem__) + + def __call__(self, *args): + return self.pmf(*args) + + def __contains__(self, other): + return other in self.set + + +#============================================= +#========= Probability Space =============== +#============================================= + + +class FinitePSpace(PSpace): + """ + A Finite Probability Space + + Represents the probabilities of a finite number of events. + """ + is_Finite = True + + def __new__(cls, domain, density): + density = {sympify(key): sympify(val) + for key, val in density.items()} + public_density = Dict(density) + + obj = PSpace.__new__(cls, domain, public_density) + obj._density = density + return obj + + def prob_of(self, elem): + elem = sympify(elem) + density = self._density + if isinstance(list(density.keys())[0], FiniteSet): + return density.get(elem, S.Zero) + return density.get(tuple(elem)[0][1], S.Zero) + + def where(self, condition): + assert all(r.symbol in self.symbols for r in random_symbols(condition)) + return ConditionalFiniteDomain(self.domain, condition) + + def compute_density(self, expr): + expr = rv_subs(expr, self.values) + d = FiniteDensity() + for elem in self.domain: + val = expr.xreplace(dict(elem)) + prob = self.prob_of(elem) + d[val] = d.get(val, S.Zero) + prob + return d + + @cacheit + def compute_cdf(self, expr): + d = self.compute_density(expr) + cum_prob = S.Zero + cdf = [] + for key in sorted(d): + prob = d[key] + cum_prob += prob + cdf.append((key, cum_prob)) + + return dict(cdf) + + @cacheit + def sorted_cdf(self, expr, python_float=False): + cdf = self.compute_cdf(expr) + items = list(cdf.items()) + sorted_items = sorted(items, key=lambda val_cumprob: val_cumprob[1]) + if python_float: + sorted_items = [(v, float(cum_prob)) + for v, cum_prob in sorted_items] + return sorted_items + + @cacheit + def compute_characteristic_function(self, expr): + d = self.compute_density(expr) + t = Dummy('t', real=True) + + return Lambda(t, sum(exp(I*k*t)*v for k,v in d.items())) + + @cacheit + def compute_moment_generating_function(self, expr): + d = self.compute_density(expr) + t = Dummy('t', real=True) + + return Lambda(t, sum(exp(k*t)*v for k,v in d.items())) + + def compute_expectation(self, expr, rvs=None, **kwargs): + rvs = rvs or self.values + expr = rv_subs(expr, rvs) + probs = [self.prob_of(elem) for elem in self.domain] + if isinstance(expr, (Logic, Relational)): + parse_domain = [tuple(elem)[0][1] for elem in self.domain] + bools = [expr.xreplace(dict(elem)) for elem in self.domain] + else: + parse_domain = [expr.xreplace(dict(elem)) for elem in self.domain] + bools = [True for elem in self.domain] + return sum(Piecewise((prob * elem, blv), (S.Zero, True)) + for prob, elem, blv in zip(probs, parse_domain, bools)) + + def compute_quantile(self, expr): + cdf = self.compute_cdf(expr) + p = Dummy('p', real=True) + set = ((nan, (p < 0) | (p > 1)),) + for key, value in cdf.items(): + set = set + ((key, p <= value), ) + return Lambda(p, Piecewise(*set)) + + def probability(self, condition): + cond_symbols = frozenset(rs.symbol for rs in random_symbols(condition)) + cond = rv_subs(condition) + if not cond_symbols.issubset(self.symbols): + raise ValueError("Cannot compare foreign random symbols, %s" + %(str(cond_symbols - self.symbols))) + if isinstance(condition, Relational) and \ + (not cond.free_symbols.issubset(self.domain.free_symbols)): + rv = condition.lhs if isinstance(condition.rhs, Symbol) else condition.rhs + return sum(Piecewise( + (self.prob_of(elem), condition.subs(rv, list(elem)[0][1])), + (S.Zero, True)) for elem in self.domain) + return sympify(sum(self.prob_of(elem) for elem in self.where(condition))) + + def conditional_space(self, condition): + domain = self.where(condition) + prob = self.probability(condition) + density = {key: val / prob + for key, val in self._density.items() if domain._test(key)} + return FinitePSpace(domain, density) + + def sample(self, size=(), library='scipy', seed=None): + """ + Internal sample method + + Returns dictionary mapping RandomSymbol to realization value. + """ + return {self.value: self.distribution.sample(size, library, seed)} + + +class SingleFinitePSpace(SinglePSpace, FinitePSpace): + """ + A single finite probability space + + Represents the probabilities of a set of random events that can be + attributed to a single variable/symbol. + + This class is implemented by many of the standard FiniteRV types such as + Die, Bernoulli, Coin, etc.... + """ + @property + def domain(self): + return SingleFiniteDomain(self.symbol, self.distribution.set) + + @property + def _is_symbolic(self): + """ + Helper property to check if the distribution + of the random variable is having symbolic + dimension. + """ + return self.distribution.is_symbolic + + @property + def distribution(self): + return self.args[1] + + def pmf(self, expr): + return self.distribution.pmf(expr) + + @property # type: ignore + @cacheit + def _density(self): + return {FiniteSet((self.symbol, val)): prob + for val, prob in self.distribution.dict.items()} + + @cacheit + def compute_characteristic_function(self, expr): + if self._is_symbolic: + d = self.compute_density(expr) + t = Dummy('t', real=True) + ki = Dummy('ki') + return Lambda(t, Sum(d(ki)*exp(I*ki*t), (ki, self.args[1].low, self.args[1].high))) + expr = rv_subs(expr, self.values) + return FinitePSpace(self.domain, self.distribution).compute_characteristic_function(expr) + + @cacheit + def compute_moment_generating_function(self, expr): + if self._is_symbolic: + d = self.compute_density(expr) + t = Dummy('t', real=True) + ki = Dummy('ki') + return Lambda(t, Sum(d(ki)*exp(ki*t), (ki, self.args[1].low, self.args[1].high))) + expr = rv_subs(expr, self.values) + return FinitePSpace(self.domain, self.distribution).compute_moment_generating_function(expr) + + def compute_quantile(self, expr): + if self._is_symbolic: + raise NotImplementedError("Computing quantile for random variables " + "with symbolic dimension because the bounds of searching the required " + "value is undetermined.") + expr = rv_subs(expr, self.values) + return FinitePSpace(self.domain, self.distribution).compute_quantile(expr) + + def compute_density(self, expr): + if self._is_symbolic: + rv = list(random_symbols(expr))[0] + k = Dummy('k', integer=True) + cond = True if not isinstance(expr, (Relational, Logic)) \ + else expr.subs(rv, k) + return Lambda(k, + Piecewise((self.pmf(k), And(k >= self.args[1].low, + k <= self.args[1].high, cond)), (S.Zero, True))) + expr = rv_subs(expr, self.values) + return FinitePSpace(self.domain, self.distribution).compute_density(expr) + + def compute_cdf(self, expr): + if self._is_symbolic: + d = self.compute_density(expr) + k = Dummy('k') + ki = Dummy('ki') + return Lambda(k, Sum(d(ki), (ki, self.args[1].low, k))) + expr = rv_subs(expr, self.values) + return FinitePSpace(self.domain, self.distribution).compute_cdf(expr) + + def compute_expectation(self, expr, rvs=None, **kwargs): + if self._is_symbolic: + rv = random_symbols(expr)[0] + k = Dummy('k', integer=True) + expr = expr.subs(rv, k) + cond = True if not isinstance(expr, (Relational, Logic)) \ + else expr + func = self.pmf(k) * k if cond != True else self.pmf(k) * expr + return Sum(Piecewise((func, cond), (S.Zero, True)), + (k, self.distribution.low, self.distribution.high)).doit() + + expr = _sympify(expr) + expr = rv_subs(expr, rvs) + return FinitePSpace(self.domain, self.distribution).compute_expectation(expr, rvs, **kwargs) + + def probability(self, condition): + if self._is_symbolic: + #TODO: Implement the mechanism for handling queries for symbolic sized distributions. + raise NotImplementedError("Currently, probability queries are not " + "supported for random variables with symbolic sized distributions.") + condition = rv_subs(condition) + return FinitePSpace(self.domain, self.distribution).probability(condition) + + def conditional_space(self, condition): + """ + This method is used for transferring the + computation to probability method because + conditional space of random variables with + symbolic dimensions is currently not possible. + """ + if self._is_symbolic: + self + domain = self.where(condition) + prob = self.probability(condition) + density = {key: val / prob + for key, val in self._density.items() if domain._test(key)} + return FinitePSpace(domain, density) + + +class ProductFinitePSpace(IndependentProductPSpace, FinitePSpace): + """ + A collection of several independent finite probability spaces + """ + @property + def domain(self): + return ProductFiniteDomain(*[space.domain for space in self.spaces]) + + @property # type: ignore + @cacheit + def _density(self): + proditer = product(*[iter(space._density.items()) + for space in self.spaces]) + d = {} + for items in proditer: + elems, probs = list(zip(*items)) + elem = sumsets(elems) + prob = Mul(*probs) + d[elem] = d.get(elem, S.Zero) + prob + return Dict(d) + + @property # type: ignore + @cacheit + def density(self): + return Dict(self._density) + + def probability(self, condition): + return FinitePSpace.probability(self, condition) + + def compute_density(self, expr): + return FinitePSpace.compute_density(self, expr) diff --git a/lib/python3.12/site-packages/sympy/stats/frv_types.py b/lib/python3.12/site-packages/sympy/stats/frv_types.py new file mode 100644 index 0000000000000000000000000000000000000000..76a099cb43a8bee7d861d3b8853aebf33cfdd24a --- /dev/null +++ b/lib/python3.12/site-packages/sympy/stats/frv_types.py @@ -0,0 +1,870 @@ +""" +Finite Discrete Random Variables - Prebuilt variable types + +Contains +======== +FiniteRV +DiscreteUniform +Die +Bernoulli +Coin +Binomial +BetaBinomial +Hypergeometric +Rademacher +IdealSoliton +RobustSoliton +""" + + +from sympy.core.cache import cacheit +from sympy.core.function import Lambda +from sympy.core.numbers import (Integer, Rational) +from sympy.core.relational import (Eq, Ge, Gt, Le, Lt) +from sympy.core.singleton import S +from sympy.core.symbol import (Dummy, Symbol) +from sympy.core.sympify import sympify +from sympy.functions.combinatorial.factorials import binomial +from sympy.functions.elementary.exponential import log +from sympy.functions.elementary.piecewise import Piecewise +from sympy.logic.boolalg import Or +from sympy.sets.contains import Contains +from sympy.sets.fancysets import Range +from sympy.sets.sets import (Intersection, Interval) +from sympy.functions.special.beta_functions import beta as beta_fn +from sympy.stats.frv import (SingleFiniteDistribution, + SingleFinitePSpace) +from sympy.stats.rv import _value_check, Density, is_random +from sympy.utilities.iterables import multiset +from sympy.utilities.misc import filldedent + + +__all__ = ['FiniteRV', +'DiscreteUniform', +'Die', +'Bernoulli', +'Coin', +'Binomial', +'BetaBinomial', +'Hypergeometric', +'Rademacher', +'IdealSoliton', +'RobustSoliton', +] + +def rv(name, cls, *args, **kwargs): + args = list(map(sympify, args)) + dist = cls(*args) + if kwargs.pop('check', True): + dist.check(*args) + pspace = SingleFinitePSpace(name, dist) + if any(is_random(arg) for arg in args): + from sympy.stats.compound_rv import CompoundPSpace, CompoundDistribution + pspace = CompoundPSpace(name, CompoundDistribution(dist)) + return pspace.value + +class FiniteDistributionHandmade(SingleFiniteDistribution): + + @property + def dict(self): + return self.args[0] + + def pmf(self, x): + x = Symbol('x') + return Lambda(x, Piecewise(*( + [(v, Eq(k, x)) for k, v in self.dict.items()] + [(S.Zero, True)]))) + + @property + def set(self): + return set(self.dict.keys()) + + @staticmethod + def check(density): + for p in density.values(): + _value_check((p >= 0, p <= 1), + "Probability at a point must be between 0 and 1.") + val = sum(density.values()) + _value_check(Eq(val, 1) != S.false, "Total Probability must be 1.") + +def FiniteRV(name, density, **kwargs): + r""" + Create a Finite Random Variable given a dict representing the density. + + Parameters + ========== + + name : Symbol + Represents name of the random variable. + density : dict + Dictionary containing the pdf of finite distribution + check : bool + If True, it will check whether the given density + integrates to 1 over the given set. If False, it + will not perform this check. Default is False. + + Examples + ======== + + >>> from sympy.stats import FiniteRV, P, E + + >>> density = {0: .1, 1: .2, 2: .3, 3: .4} + >>> X = FiniteRV('X', density) + + >>> E(X) + 2.00000000000000 + >>> P(X >= 2) + 0.700000000000000 + + Returns + ======= + + RandomSymbol + + """ + # have a default of False while `rv` should have a default of True + kwargs['check'] = kwargs.pop('check', False) + return rv(name, FiniteDistributionHandmade, density, **kwargs) + +class DiscreteUniformDistribution(SingleFiniteDistribution): + + @staticmethod + def check(*args): + # not using _value_check since there is a + # suggestion for the user + if len(set(args)) != len(args): + weights = multiset(args) + n = Integer(len(args)) + for k in weights: + weights[k] /= n + raise ValueError(filldedent(""" + Repeated args detected but set expected. For a + distribution having different weights for each + item use the following:""") + ( + '\nS("FiniteRV(%s, %s)")' % ("'X'", weights))) + + @property + def p(self): + return Rational(1, len(self.args)) + + @property # type: ignore + @cacheit + def dict(self): + return dict.fromkeys(self.set, self.p) + + @property + def set(self): + return set(self.args) + + def pmf(self, x): + if x in self.args: + return self.p + else: + return S.Zero + + +def DiscreteUniform(name, items): + r""" + Create a Finite Random Variable representing a uniform distribution over + the input set. + + Parameters + ========== + + items : list/tuple + Items over which Uniform distribution is to be made + + Examples + ======== + + >>> from sympy.stats import DiscreteUniform, density + >>> from sympy import symbols + + >>> X = DiscreteUniform('X', symbols('a b c')) # equally likely over a, b, c + >>> density(X).dict + {a: 1/3, b: 1/3, c: 1/3} + + >>> Y = DiscreteUniform('Y', list(range(5))) # distribution over a range + >>> density(Y).dict + {0: 1/5, 1: 1/5, 2: 1/5, 3: 1/5, 4: 1/5} + + Returns + ======= + + RandomSymbol + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Discrete_uniform_distribution + .. [2] https://mathworld.wolfram.com/DiscreteUniformDistribution.html + + """ + return rv(name, DiscreteUniformDistribution, *items) + + +class DieDistribution(SingleFiniteDistribution): + _argnames = ('sides',) + + @staticmethod + def check(sides): + _value_check((sides.is_positive, sides.is_integer), + "number of sides must be a positive integer.") + + @property + def is_symbolic(self): + return not self.sides.is_number + + @property + def high(self): + return self.sides + + @property + def low(self): + return S.One + + @property + def set(self): + if self.is_symbolic: + return Intersection(S.Naturals0, Interval(0, self.sides)) + return set(map(Integer, range(1, self.sides + 1))) + + def pmf(self, x): + x = sympify(x) + if not (x.is_number or x.is_Symbol or is_random(x)): + raise ValueError("'x' expected as an argument of type 'number', 'Symbol', or " + "'RandomSymbol' not %s" % (type(x))) + cond = Ge(x, 1) & Le(x, self.sides) & Contains(x, S.Integers) + return Piecewise((S.One/self.sides, cond), (S.Zero, True)) + +def Die(name, sides=6): + r""" + Create a Finite Random Variable representing a fair die. + + Parameters + ========== + + sides : Integer + Represents the number of sides of the Die, by default is 6 + + Examples + ======== + + >>> from sympy.stats import Die, density + >>> from sympy import Symbol + + >>> D6 = Die('D6', 6) # Six sided Die + >>> density(D6).dict + {1: 1/6, 2: 1/6, 3: 1/6, 4: 1/6, 5: 1/6, 6: 1/6} + + >>> D4 = Die('D4', 4) # Four sided Die + >>> density(D4).dict + {1: 1/4, 2: 1/4, 3: 1/4, 4: 1/4} + + >>> n = Symbol('n', positive=True, integer=True) + >>> Dn = Die('Dn', n) # n sided Die + >>> density(Dn).dict + Density(DieDistribution(n)) + >>> density(Dn).dict.subs(n, 4).doit() + {1: 1/4, 2: 1/4, 3: 1/4, 4: 1/4} + + Returns + ======= + + RandomSymbol + """ + + return rv(name, DieDistribution, sides) + + +class BernoulliDistribution(SingleFiniteDistribution): + _argnames = ('p', 'succ', 'fail') + + @staticmethod + def check(p, succ, fail): + _value_check((p >= 0, p <= 1), + "p should be in range [0, 1].") + + @property + def set(self): + return {self.succ, self.fail} + + def pmf(self, x): + if isinstance(self.succ, Symbol) and isinstance(self.fail, Symbol): + return Piecewise((self.p, x == self.succ), + (1 - self.p, x == self.fail), + (S.Zero, True)) + return Piecewise((self.p, Eq(x, self.succ)), + (1 - self.p, Eq(x, self.fail)), + (S.Zero, True)) + + +def Bernoulli(name, p, succ=1, fail=0): + r""" + Create a Finite Random Variable representing a Bernoulli process. + + Parameters + ========== + + p : Rational number between 0 and 1 + Represents probability of success + succ : Integer/symbol/string + Represents event of success + fail : Integer/symbol/string + Represents event of failure + + Examples + ======== + + >>> from sympy.stats import Bernoulli, density + >>> from sympy import S + + >>> X = Bernoulli('X', S(3)/4) # 1-0 Bernoulli variable, probability = 3/4 + >>> density(X).dict + {0: 1/4, 1: 3/4} + + >>> X = Bernoulli('X', S.Half, 'Heads', 'Tails') # A fair coin toss + >>> density(X).dict + {Heads: 1/2, Tails: 1/2} + + Returns + ======= + + RandomSymbol + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Bernoulli_distribution + .. [2] https://mathworld.wolfram.com/BernoulliDistribution.html + + """ + + return rv(name, BernoulliDistribution, p, succ, fail) + + +def Coin(name, p=S.Half): + r""" + Create a Finite Random Variable representing a Coin toss. + + Parameters + ========== + + p : Rational Number between 0 and 1 + Represents probability of getting "Heads", by default is Half + + Examples + ======== + + >>> from sympy.stats import Coin, density + >>> from sympy import Rational + + >>> C = Coin('C') # A fair coin toss + >>> density(C).dict + {H: 1/2, T: 1/2} + + >>> C2 = Coin('C2', Rational(3, 5)) # An unfair coin + >>> density(C2).dict + {H: 3/5, T: 2/5} + + Returns + ======= + + RandomSymbol + + See Also + ======== + + sympy.stats.Binomial + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Coin_flipping + + """ + return rv(name, BernoulliDistribution, p, 'H', 'T') + + +class BinomialDistribution(SingleFiniteDistribution): + _argnames = ('n', 'p', 'succ', 'fail') + + @staticmethod + def check(n, p, succ, fail): + _value_check((n.is_integer, n.is_nonnegative), + "'n' must be nonnegative integer.") + _value_check((p <= 1, p >= 0), + "p should be in range [0, 1].") + + @property + def high(self): + return self.n + + @property + def low(self): + return S.Zero + + @property + def is_symbolic(self): + return not self.n.is_number + + @property + def set(self): + if self.is_symbolic: + return Intersection(S.Naturals0, Interval(0, self.n)) + return set(self.dict.keys()) + + def pmf(self, x): + n, p = self.n, self.p + x = sympify(x) + if not (x.is_number or x.is_Symbol or is_random(x)): + raise ValueError("'x' expected as an argument of type 'number', 'Symbol', or " + "'RandomSymbol' not %s" % (type(x))) + cond = Ge(x, 0) & Le(x, n) & Contains(x, S.Integers) + return Piecewise((binomial(n, x) * p**x * (1 - p)**(n - x), cond), (S.Zero, True)) + + @property # type: ignore + @cacheit + def dict(self): + if self.is_symbolic: + return Density(self) + return {k*self.succ + (self.n-k)*self.fail: self.pmf(k) + for k in range(0, self.n + 1)} + + +def Binomial(name, n, p, succ=1, fail=0): + r""" + Create a Finite Random Variable representing a binomial distribution. + + Parameters + ========== + + n : Positive Integer + Represents number of trials + p : Rational Number between 0 and 1 + Represents probability of success + succ : Integer/symbol/string + Represents event of success, by default is 1 + fail : Integer/symbol/string + Represents event of failure, by default is 0 + + Examples + ======== + + >>> from sympy.stats import Binomial, density + >>> from sympy import S, Symbol + + >>> X = Binomial('X', 4, S.Half) # Four "coin flips" + >>> density(X).dict + {0: 1/16, 1: 1/4, 2: 3/8, 3: 1/4, 4: 1/16} + + >>> n = Symbol('n', positive=True, integer=True) + >>> p = Symbol('p', positive=True) + >>> X = Binomial('X', n, S.Half) # n "coin flips" + >>> density(X).dict + Density(BinomialDistribution(n, 1/2, 1, 0)) + >>> density(X).dict.subs(n, 4).doit() + {0: 1/16, 1: 1/4, 2: 3/8, 3: 1/4, 4: 1/16} + + Returns + ======= + + RandomSymbol + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Binomial_distribution + .. [2] https://mathworld.wolfram.com/BinomialDistribution.html + + """ + + return rv(name, BinomialDistribution, n, p, succ, fail) + +#------------------------------------------------------------------------------- +# Beta-binomial distribution ---------------------------------------------------------- + +class BetaBinomialDistribution(SingleFiniteDistribution): + _argnames = ('n', 'alpha', 'beta') + + @staticmethod + def check(n, alpha, beta): + _value_check((n.is_integer, n.is_nonnegative), + "'n' must be nonnegative integer. n = %s." % str(n)) + _value_check((alpha > 0), + "'alpha' must be: alpha > 0 . alpha = %s" % str(alpha)) + _value_check((beta > 0), + "'beta' must be: beta > 0 . beta = %s" % str(beta)) + + @property + def high(self): + return self.n + + @property + def low(self): + return S.Zero + + @property + def is_symbolic(self): + return not self.n.is_number + + @property + def set(self): + if self.is_symbolic: + return Intersection(S.Naturals0, Interval(0, self.n)) + return set(map(Integer, range(self.n + 1))) + + def pmf(self, k): + n, a, b = self.n, self.alpha, self.beta + return binomial(n, k) * beta_fn(k + a, n - k + b) / beta_fn(a, b) + + +def BetaBinomial(name, n, alpha, beta): + r""" + Create a Finite Random Variable representing a Beta-binomial distribution. + + Parameters + ========== + + n : Positive Integer + Represents number of trials + alpha : Real positive number + beta : Real positive number + + Examples + ======== + + >>> from sympy.stats import BetaBinomial, density + + >>> X = BetaBinomial('X', 2, 1, 1) + >>> density(X).dict + {0: 1/3, 1: 2*beta(2, 2), 2: 1/3} + + Returns + ======= + + RandomSymbol + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Beta-binomial_distribution + .. [2] https://mathworld.wolfram.com/BetaBinomialDistribution.html + + """ + + return rv(name, BetaBinomialDistribution, n, alpha, beta) + + +class HypergeometricDistribution(SingleFiniteDistribution): + _argnames = ('N', 'm', 'n') + + @staticmethod + def check(n, N, m): + _value_check((N.is_integer, N.is_nonnegative), + "'N' must be nonnegative integer. N = %s." % str(N)) + _value_check((n.is_integer, n.is_nonnegative), + "'n' must be nonnegative integer. n = %s." % str(n)) + _value_check((m.is_integer, m.is_nonnegative), + "'m' must be nonnegative integer. m = %s." % str(m)) + + @property + def is_symbolic(self): + return not all(x.is_number for x in (self.N, self.m, self.n)) + + @property + def high(self): + return Piecewise((self.n, Lt(self.n, self.m) != False), (self.m, True)) + + @property + def low(self): + return Piecewise((0, Gt(0, self.n + self.m - self.N) != False), (self.n + self.m - self.N, True)) + + @property + def set(self): + N, m, n = self.N, self.m, self.n + if self.is_symbolic: + return Intersection(S.Naturals0, Interval(self.low, self.high)) + return set(range(max(0, n + m - N), min(n, m) + 1)) + + def pmf(self, k): + N, m, n = self.N, self.m, self.n + return S(binomial(m, k) * binomial(N - m, n - k))/binomial(N, n) + + +def Hypergeometric(name, N, m, n): + r""" + Create a Finite Random Variable representing a hypergeometric distribution. + + Parameters + ========== + + N : Positive Integer + Represents finite population of size N. + m : Positive Integer + Represents number of trials with required feature. + n : Positive Integer + Represents numbers of draws. + + + Examples + ======== + + >>> from sympy.stats import Hypergeometric, density + + >>> X = Hypergeometric('X', 10, 5, 3) # 10 marbles, 5 white (success), 3 draws + >>> density(X).dict + {0: 1/12, 1: 5/12, 2: 5/12, 3: 1/12} + + Returns + ======= + + RandomSymbol + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Hypergeometric_distribution + .. [2] https://mathworld.wolfram.com/HypergeometricDistribution.html + + """ + return rv(name, HypergeometricDistribution, N, m, n) + + +class RademacherDistribution(SingleFiniteDistribution): + + @property + def set(self): + return {-1, 1} + + @property + def pmf(self): + k = Dummy('k') + return Lambda(k, Piecewise((S.Half, Or(Eq(k, -1), Eq(k, 1))), (S.Zero, True))) + +def Rademacher(name): + r""" + Create a Finite Random Variable representing a Rademacher distribution. + + Examples + ======== + + >>> from sympy.stats import Rademacher, density + + >>> X = Rademacher('X') + >>> density(X).dict + {-1: 1/2, 1: 1/2} + + Returns + ======= + + RandomSymbol + + See Also + ======== + + sympy.stats.Bernoulli + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Rademacher_distribution + + """ + return rv(name, RademacherDistribution) + +class IdealSolitonDistribution(SingleFiniteDistribution): + _argnames = ('k',) + + @staticmethod + def check(k): + _value_check(k.is_integer and k.is_positive, + "'k' must be a positive integer.") + + @property + def low(self): + return S.One + + @property + def high(self): + return self.k + + @property + def set(self): + return set(map(Integer, range(1, self.k + 1))) + + @property # type: ignore + @cacheit + def dict(self): + if self.k.is_Symbol: + return Density(self) + d = {1: Rational(1, self.k)} + d.update({i: Rational(1, i*(i - 1)) for i in range(2, self.k + 1)}) + return d + + def pmf(self, x): + x = sympify(x) + if not (x.is_number or x.is_Symbol or is_random(x)): + raise ValueError("'x' expected as an argument of type 'number', 'Symbol', or " + "'RandomSymbol' not %s" % (type(x))) + cond1 = Eq(x, 1) & x.is_integer + cond2 = Ge(x, 1) & Le(x, self.k) & x.is_integer + return Piecewise((1/self.k, cond1), (1/(x*(x - 1)), cond2), (S.Zero, True)) + +def IdealSoliton(name, k): + r""" + Create a Finite Random Variable of Ideal Soliton Distribution + + Parameters + ========== + + k : Positive Integer + Represents the number of input symbols in an LT (Luby Transform) code. + + Examples + ======== + + >>> from sympy.stats import IdealSoliton, density, P, E + >>> sol = IdealSoliton('sol', 5) + >>> density(sol).dict + {1: 1/5, 2: 1/2, 3: 1/6, 4: 1/12, 5: 1/20} + >>> density(sol).set + {1, 2, 3, 4, 5} + + >>> from sympy import Symbol + >>> k = Symbol('k', positive=True, integer=True) + >>> sol = IdealSoliton('sol', k) + >>> density(sol).dict + Density(IdealSolitonDistribution(k)) + >>> density(sol).dict.subs(k, 10).doit() + {1: 1/10, 2: 1/2, 3: 1/6, 4: 1/12, 5: 1/20, 6: 1/30, 7: 1/42, 8: 1/56, 9: 1/72, 10: 1/90} + + >>> E(sol.subs(k, 10)) + 7381/2520 + + >>> P(sol.subs(k, 4) > 2) + 1/4 + + Returns + ======= + + RandomSymbol + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Soliton_distribution#Ideal_distribution + .. [2] https://pages.cs.wisc.edu/~suman/courses/740/papers/luby02lt.pdf + + """ + return rv(name, IdealSolitonDistribution, k) + +class RobustSolitonDistribution(SingleFiniteDistribution): + _argnames= ('k', 'delta', 'c') + + @staticmethod + def check(k, delta, c): + _value_check(k.is_integer and k.is_positive, + "'k' must be a positive integer") + _value_check(Gt(delta, 0) and Le(delta, 1), + "'delta' must be a real number in the interval (0,1)") + _value_check(c.is_positive, + "'c' must be a positive real number.") + + @property + def R(self): + return self.c * log(self.k/self.delta) * self.k**0.5 + + @property + def Z(self): + z = 0 + for i in Range(1, round(self.k/self.R)): + z += (1/i) + z += log(self.R/self.delta) + return 1 + z * self.R/self.k + + @property + def low(self): + return S.One + + @property + def high(self): + return self.k + + @property + def set(self): + return set(map(Integer, range(1, self.k + 1))) + + @property + def is_symbolic(self): + return not (self.k.is_number and self.c.is_number and self.delta.is_number) + + def pmf(self, x): + x = sympify(x) + if not (x.is_number or x.is_Symbol or is_random(x)): + raise ValueError("'x' expected as an argument of type 'number', 'Symbol', or " + "'RandomSymbol' not %s" % (type(x))) + + cond1 = Eq(x, 1) & x.is_integer + cond2 = Ge(x, 1) & Le(x, self.k) & x.is_integer + rho = Piecewise((Rational(1, self.k), cond1), (Rational(1, x*(x-1)), cond2), (S.Zero, True)) + + cond1 = Ge(x, 1) & Le(x, round(self.k/self.R)-1) + cond2 = Eq(x, round(self.k/self.R)) + tau = Piecewise((self.R/(self.k * x), cond1), (self.R * log(self.R/self.delta)/self.k, cond2), (S.Zero, True)) + + return (rho + tau)/self.Z + +def RobustSoliton(name, k, delta, c): + r''' + Create a Finite Random Variable of Robust Soliton Distribution + + Parameters + ========== + + k : Positive Integer + Represents the number of input symbols in an LT (Luby Transform) code. + delta : Positive Rational Number + Represents the failure probability. Must be in the interval (0,1). + c : Positive Rational Number + Constant of proportionality. Values close to 1 are recommended + + Examples + ======== + + >>> from sympy.stats import RobustSoliton, density, P, E + >>> robSol = RobustSoliton('robSol', 5, 0.5, 0.01) + >>> density(robSol).dict + {1: 0.204253668152708, 2: 0.490631107897393, 3: 0.165210624506162, 4: 0.0834387731899302, 5: 0.0505633404760675} + >>> density(robSol).set + {1, 2, 3, 4, 5} + + >>> from sympy import Symbol + >>> k = Symbol('k', positive=True, integer=True) + >>> c = Symbol('c', positive=True) + >>> robSol = RobustSoliton('robSol', k, 0.5, c) + >>> density(robSol).dict + Density(RobustSolitonDistribution(k, 0.5, c)) + >>> density(robSol).dict.subs(k, 10).subs(c, 0.03).doit() + {1: 0.116641095387194, 2: 0.467045731687165, 3: 0.159984123349381, 4: 0.0821431680681869, 5: 0.0505765646770100, + 6: 0.0345781523420719, 7: 0.0253132820710503, 8: 0.0194459129233227, 9: 0.0154831166726115, 10: 0.0126733075238887} + + >>> E(robSol.subs(k, 10).subs(c, 0.05)) + 2.91358846104106 + + >>> P(robSol.subs(k, 4).subs(c, 0.1) > 2) + 0.243650614389834 + + Returns + ======= + + RandomSymbol + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Soliton_distribution#Robust_distribution + .. [2] https://www.inference.org.uk/mackay/itprnn/ps/588.596.pdf + .. [3] https://pages.cs.wisc.edu/~suman/courses/740/papers/luby02lt.pdf + + ''' + return rv(name, RobustSolitonDistribution, k, delta, c) diff --git a/lib/python3.12/site-packages/sympy/stats/joint_rv.py b/lib/python3.12/site-packages/sympy/stats/joint_rv.py new file mode 100644 index 0000000000000000000000000000000000000000..d147942f08b998e167b246628360fa27fc8ef348 --- /dev/null +++ b/lib/python3.12/site-packages/sympy/stats/joint_rv.py @@ -0,0 +1,426 @@ +""" +Joint Random Variables Module + +See Also +======== +sympy.stats.rv +sympy.stats.frv +sympy.stats.crv +sympy.stats.drv +""" +from math import prod + +from sympy.core.basic import Basic +from sympy.core.function import Lambda +from sympy.core.singleton import S +from sympy.core.symbol import (Dummy, Symbol) +from sympy.core.sympify import sympify +from sympy.sets.sets import ProductSet +from sympy.tensor.indexed import Indexed +from sympy.concrete.products import Product +from sympy.concrete.summations import Sum, summation +from sympy.core.containers import Tuple +from sympy.integrals.integrals import Integral, integrate +from sympy.matrices import ImmutableMatrix, matrix2numpy, list2numpy +from sympy.stats.crv import SingleContinuousDistribution, SingleContinuousPSpace +from sympy.stats.drv import SingleDiscreteDistribution, SingleDiscretePSpace +from sympy.stats.rv import (ProductPSpace, NamedArgsMixin, Distribution, + ProductDomain, RandomSymbol, random_symbols, + SingleDomain, _symbol_converter) +from sympy.utilities.iterables import iterable +from sympy.utilities.misc import filldedent +from sympy.external import import_module + +# __all__ = ['marginal_distribution'] + +class JointPSpace(ProductPSpace): + """ + Represents a joint probability space. Represented using symbols for + each component and a distribution. + """ + def __new__(cls, sym, dist): + if isinstance(dist, SingleContinuousDistribution): + return SingleContinuousPSpace(sym, dist) + if isinstance(dist, SingleDiscreteDistribution): + return SingleDiscretePSpace(sym, dist) + sym = _symbol_converter(sym) + return Basic.__new__(cls, sym, dist) + + @property + def set(self): + return self.domain.set + + @property + def symbol(self): + return self.args[0] + + @property + def distribution(self): + return self.args[1] + + @property + def value(self): + return JointRandomSymbol(self.symbol, self) + + @property + def component_count(self): + _set = self.distribution.set + if isinstance(_set, ProductSet): + return S(len(_set.args)) + elif isinstance(_set, Product): + return _set.limits[0][-1] + return S.One + + @property + def pdf(self): + sym = [Indexed(self.symbol, i) for i in range(self.component_count)] + return self.distribution(*sym) + + @property + def domain(self): + rvs = random_symbols(self.distribution) + if not rvs: + return SingleDomain(self.symbol, self.distribution.set) + return ProductDomain(*[rv.pspace.domain for rv in rvs]) + + def component_domain(self, index): + return self.set.args[index] + + def marginal_distribution(self, *indices): + count = self.component_count + if count.atoms(Symbol): + raise ValueError("Marginal distributions cannot be computed " + "for symbolic dimensions. It is a work under progress.") + orig = [Indexed(self.symbol, i) for i in range(count)] + all_syms = [Symbol(str(i)) for i in orig] + replace_dict = dict(zip(all_syms, orig)) + sym = tuple(Symbol(str(Indexed(self.symbol, i))) for i in indices) + limits = [[i,] for i in all_syms if i not in sym] + index = 0 + for i in range(count): + if i not in indices: + limits[index].append(self.distribution.set.args[i]) + limits[index] = tuple(limits[index]) + index += 1 + if self.distribution.is_Continuous: + f = Lambda(sym, integrate(self.distribution(*all_syms), *limits)) + elif self.distribution.is_Discrete: + f = Lambda(sym, summation(self.distribution(*all_syms), *limits)) + return f.xreplace(replace_dict) + + def compute_expectation(self, expr, rvs=None, evaluate=False, **kwargs): + syms = tuple(self.value[i] for i in range(self.component_count)) + rvs = rvs or syms + if not any(i in rvs for i in syms): + return expr + expr = expr*self.pdf + for rv in rvs: + if isinstance(rv, Indexed): + expr = expr.xreplace({rv: Indexed(str(rv.base), rv.args[1])}) + elif isinstance(rv, RandomSymbol): + expr = expr.xreplace({rv: rv.symbol}) + if self.value in random_symbols(expr): + raise NotImplementedError(filldedent(''' + Expectations of expression with unindexed joint random symbols + cannot be calculated yet.''')) + limits = tuple((Indexed(str(rv.base),rv.args[1]), + self.distribution.set.args[rv.args[1]]) for rv in syms) + return Integral(expr, *limits) + + def where(self, condition): + raise NotImplementedError() + + def compute_density(self, expr): + raise NotImplementedError() + + def sample(self, size=(), library='scipy', seed=None): + """ + Internal sample method + + Returns dictionary mapping RandomSymbol to realization value. + """ + return {RandomSymbol(self.symbol, self): self.distribution.sample(size, + library=library, seed=seed)} + + def probability(self, condition): + raise NotImplementedError() + + +class SampleJointScipy: + """Returns the sample from scipy of the given distribution""" + def __new__(cls, dist, size, seed=None): + return cls._sample_scipy(dist, size, seed) + + @classmethod + def _sample_scipy(cls, dist, size, seed): + """Sample from SciPy.""" + + import numpy + if seed is None or isinstance(seed, int): + rand_state = numpy.random.default_rng(seed=seed) + else: + rand_state = seed + from scipy import stats as scipy_stats + scipy_rv_map = { + 'MultivariateNormalDistribution': lambda dist, size: scipy_stats.multivariate_normal.rvs( + mean=matrix2numpy(dist.mu).flatten(), + cov=matrix2numpy(dist.sigma), size=size, random_state=rand_state), + 'MultivariateBetaDistribution': lambda dist, size: scipy_stats.dirichlet.rvs( + alpha=list2numpy(dist.alpha, float).flatten(), size=size, random_state=rand_state), + 'MultinomialDistribution': lambda dist, size: scipy_stats.multinomial.rvs( + n=int(dist.n), p=list2numpy(dist.p, float).flatten(), size=size, random_state=rand_state) + } + + sample_shape = { + 'MultivariateNormalDistribution': lambda dist: matrix2numpy(dist.mu).flatten().shape, + 'MultivariateBetaDistribution': lambda dist: list2numpy(dist.alpha).flatten().shape, + 'MultinomialDistribution': lambda dist: list2numpy(dist.p).flatten().shape + } + + dist_list = scipy_rv_map.keys() + + if dist.__class__.__name__ not in dist_list: + return None + + samples = scipy_rv_map[dist.__class__.__name__](dist, size) + return samples.reshape(size + sample_shape[dist.__class__.__name__](dist)) + +class SampleJointNumpy: + """Returns the sample from numpy of the given distribution""" + + def __new__(cls, dist, size, seed=None): + return cls._sample_numpy(dist, size, seed) + + @classmethod + def _sample_numpy(cls, dist, size, seed): + """Sample from NumPy.""" + + import numpy + if seed is None or isinstance(seed, int): + rand_state = numpy.random.default_rng(seed=seed) + else: + rand_state = seed + numpy_rv_map = { + 'MultivariateNormalDistribution': lambda dist, size: rand_state.multivariate_normal( + mean=matrix2numpy(dist.mu, float).flatten(), + cov=matrix2numpy(dist.sigma, float), size=size), + 'MultivariateBetaDistribution': lambda dist, size: rand_state.dirichlet( + alpha=list2numpy(dist.alpha, float).flatten(), size=size), + 'MultinomialDistribution': lambda dist, size: rand_state.multinomial( + n=int(dist.n), pvals=list2numpy(dist.p, float).flatten(), size=size) + } + + sample_shape = { + 'MultivariateNormalDistribution': lambda dist: matrix2numpy(dist.mu).flatten().shape, + 'MultivariateBetaDistribution': lambda dist: list2numpy(dist.alpha).flatten().shape, + 'MultinomialDistribution': lambda dist: list2numpy(dist.p).flatten().shape + } + + dist_list = numpy_rv_map.keys() + + if dist.__class__.__name__ not in dist_list: + return None + + samples = numpy_rv_map[dist.__class__.__name__](dist, prod(size)) + return samples.reshape(size + sample_shape[dist.__class__.__name__](dist)) + +class SampleJointPymc: + """Returns the sample from pymc of the given distribution""" + + def __new__(cls, dist, size, seed=None): + return cls._sample_pymc(dist, size, seed) + + @classmethod + def _sample_pymc(cls, dist, size, seed): + """Sample from PyMC.""" + + try: + import pymc + except ImportError: + import pymc3 as pymc + pymc_rv_map = { + 'MultivariateNormalDistribution': lambda dist: + pymc.MvNormal('X', mu=matrix2numpy(dist.mu, float).flatten(), + cov=matrix2numpy(dist.sigma, float), shape=(1, dist.mu.shape[0])), + 'MultivariateBetaDistribution': lambda dist: + pymc.Dirichlet('X', a=list2numpy(dist.alpha, float).flatten()), + 'MultinomialDistribution': lambda dist: + pymc.Multinomial('X', n=int(dist.n), + p=list2numpy(dist.p, float).flatten(), shape=(1, len(dist.p))) + } + + sample_shape = { + 'MultivariateNormalDistribution': lambda dist: matrix2numpy(dist.mu).flatten().shape, + 'MultivariateBetaDistribution': lambda dist: list2numpy(dist.alpha).flatten().shape, + 'MultinomialDistribution': lambda dist: list2numpy(dist.p).flatten().shape + } + + dist_list = pymc_rv_map.keys() + + if dist.__class__.__name__ not in dist_list: + return None + + import logging + logging.getLogger("pymc3").setLevel(logging.ERROR) + with pymc.Model(): + pymc_rv_map[dist.__class__.__name__](dist) + samples = pymc.sample(draws=prod(size), chains=1, progressbar=False, random_seed=seed, return_inferencedata=False, compute_convergence_checks=False)[:]['X'] + return samples.reshape(size + sample_shape[dist.__class__.__name__](dist)) + + +_get_sample_class_jrv = { + 'scipy': SampleJointScipy, + 'pymc3': SampleJointPymc, + 'pymc': SampleJointPymc, + 'numpy': SampleJointNumpy +} + +class JointDistribution(Distribution, NamedArgsMixin): + """ + Represented by the random variables part of the joint distribution. + Contains methods for PDF, CDF, sampling, marginal densities, etc. + """ + + _argnames = ('pdf', ) + + def __new__(cls, *args): + args = list(map(sympify, args)) + for i in range(len(args)): + if isinstance(args[i], list): + args[i] = ImmutableMatrix(args[i]) + return Basic.__new__(cls, *args) + + @property + def domain(self): + return ProductDomain(self.symbols) + + @property + def pdf(self): + return self.density.args[1] + + def cdf(self, other): + if not isinstance(other, dict): + raise ValueError("%s should be of type dict, got %s"%(other, type(other))) + rvs = other.keys() + _set = self.domain.set.sets + expr = self.pdf(tuple(i.args[0] for i in self.symbols)) + for i in range(len(other)): + if rvs[i].is_Continuous: + density = Integral(expr, (rvs[i], _set[i].inf, + other[rvs[i]])) + elif rvs[i].is_Discrete: + density = Sum(expr, (rvs[i], _set[i].inf, + other[rvs[i]])) + return density + + def sample(self, size=(), library='scipy', seed=None): + """ A random realization from the distribution """ + + libraries = ('scipy', 'numpy', 'pymc3', 'pymc') + if library not in libraries: + raise NotImplementedError("Sampling from %s is not supported yet." + % str(library)) + if not import_module(library): + raise ValueError("Failed to import %s" % library) + + samps = _get_sample_class_jrv[library](self, size, seed=seed) + + if samps is not None: + return samps + raise NotImplementedError( + "Sampling for %s is not currently implemented from %s" + % (self.__class__.__name__, library) + ) + + def __call__(self, *args): + return self.pdf(*args) + +class JointRandomSymbol(RandomSymbol): + """ + Representation of random symbols with joint probability distributions + to allow indexing." + """ + def __getitem__(self, key): + if isinstance(self.pspace, JointPSpace): + if (self.pspace.component_count <= key) == True: + raise ValueError("Index keys for %s can only up to %s." % + (self.name, self.pspace.component_count - 1)) + return Indexed(self, key) + + + +class MarginalDistribution(Distribution): + """ + Represents the marginal distribution of a joint probability space. + + Initialised using a probability distribution and random variables(or + their indexed components) which should be a part of the resultant + distribution. + """ + + def __new__(cls, dist, *rvs): + if len(rvs) == 1 and iterable(rvs[0]): + rvs = tuple(rvs[0]) + if not all(isinstance(rv, (Indexed, RandomSymbol)) for rv in rvs): + raise ValueError(filldedent('''Marginal distribution can be + intitialised only in terms of random variables or indexed random + variables''')) + rvs = Tuple.fromiter(rv for rv in rvs) + if not isinstance(dist, JointDistribution) and len(random_symbols(dist)) == 0: + return dist + return Basic.__new__(cls, dist, rvs) + + def check(self): + pass + + @property + def set(self): + rvs = [i for i in self.args[1] if isinstance(i, RandomSymbol)] + return ProductSet(*[rv.pspace.set for rv in rvs]) + + @property + def symbols(self): + rvs = self.args[1] + return {rv.pspace.symbol for rv in rvs} + + def pdf(self, *x): + expr, rvs = self.args[0], self.args[1] + marginalise_out = [i for i in random_symbols(expr) if i not in rvs] + if isinstance(expr, JointDistribution): + count = len(expr.domain.args) + x = Dummy('x', real=True) + syms = tuple(Indexed(x, i) for i in count) + expr = expr.pdf(syms) + else: + syms = tuple(rv.pspace.symbol if isinstance(rv, RandomSymbol) else rv.args[0] for rv in rvs) + return Lambda(syms, self.compute_pdf(expr, marginalise_out))(*x) + + def compute_pdf(self, expr, rvs): + for rv in rvs: + lpdf = 1 + if isinstance(rv, RandomSymbol): + lpdf = rv.pspace.pdf + expr = self.marginalise_out(expr*lpdf, rv) + return expr + + def marginalise_out(self, expr, rv): + from sympy.concrete.summations import Sum + if isinstance(rv, RandomSymbol): + dom = rv.pspace.set + elif isinstance(rv, Indexed): + dom = rv.base.component_domain( + rv.pspace.component_domain(rv.args[1])) + expr = expr.xreplace({rv: rv.pspace.symbol}) + if rv.pspace.is_Continuous: + #TODO: Modify to support integration + #for all kinds of sets. + expr = Integral(expr, (rv.pspace.symbol, dom)) + elif rv.pspace.is_Discrete: + #incorporate this into `Sum`/`summation` + if dom in (S.Integers, S.Naturals, S.Naturals0): + dom = (dom.inf, dom.sup) + expr = Sum(expr, (rv.pspace.symbol, dom)) + return expr + + def __call__(self, *args): + return self.pdf(*args) diff --git a/lib/python3.12/site-packages/sympy/stats/joint_rv_types.py b/lib/python3.12/site-packages/sympy/stats/joint_rv_types.py new file mode 100644 index 0000000000000000000000000000000000000000..6cee9f9aa30897593ffb7c7b930a55a38f0c518a --- /dev/null +++ b/lib/python3.12/site-packages/sympy/stats/joint_rv_types.py @@ -0,0 +1,945 @@ +from sympy.concrete.products import Product +from sympy.concrete.summations import Sum +from sympy.core.add import Add +from sympy.core.function import Lambda +from sympy.core.mul import Mul +from sympy.core.numbers import (Integer, Rational, pi) +from sympy.core.power import Pow +from sympy.core.relational import Eq +from sympy.core.singleton import S +from sympy.core.symbol import (Symbol, symbols) +from sympy.core.sympify import sympify +from sympy.functions.combinatorial.factorials import (rf, factorial) +from sympy.functions.elementary.exponential import exp +from sympy.functions.elementary.miscellaneous import sqrt +from sympy.functions.elementary.piecewise import Piecewise +from sympy.functions.special.bessel import besselk +from sympy.functions.special.gamma_functions import gamma +from sympy.matrices.dense import (Matrix, ones) +from sympy.sets.fancysets import Range +from sympy.sets.sets import (Intersection, Interval) +from sympy.tensor.indexed import (Indexed, IndexedBase) +from sympy.matrices import ImmutableMatrix, MatrixSymbol +from sympy.matrices.expressions.determinant import det +from sympy.matrices.expressions.matexpr import MatrixElement +from sympy.stats.joint_rv import JointDistribution, JointPSpace, MarginalDistribution +from sympy.stats.rv import _value_check, random_symbols + +__all__ = ['JointRV', +'MultivariateNormal', +'MultivariateLaplace', +'Dirichlet', +'GeneralizedMultivariateLogGamma', +'GeneralizedMultivariateLogGammaOmega', +'Multinomial', +'MultivariateBeta', +'MultivariateEwens', +'MultivariateT', +'NegativeMultinomial', +'NormalGamma' +] + +def multivariate_rv(cls, sym, *args): + args = list(map(sympify, args)) + dist = cls(*args) + args = dist.args + dist.check(*args) + return JointPSpace(sym, dist).value + + +def marginal_distribution(rv, *indices): + """ + Marginal distribution function of a joint random variable. + + Parameters + ========== + + rv : A random variable with a joint probability distribution. + indices : Component indices or the indexed random symbol + for which the joint distribution is to be calculated + + Returns + ======= + + A Lambda expression in `sym`. + + Examples + ======== + + >>> from sympy.stats import MultivariateNormal, marginal_distribution + >>> m = MultivariateNormal('X', [1, 2], [[2, 1], [1, 2]]) + >>> marginal_distribution(m, m[0])(1) + 1/(2*sqrt(pi)) + + """ + indices = list(indices) + for i in range(len(indices)): + if isinstance(indices[i], Indexed): + indices[i] = indices[i].args[1] + prob_space = rv.pspace + if not indices: + raise ValueError( + "At least one component for marginal density is needed.") + if hasattr(prob_space.distribution, '_marginal_distribution'): + return prob_space.distribution._marginal_distribution(indices, rv.symbol) + return prob_space.marginal_distribution(*indices) + + +class JointDistributionHandmade(JointDistribution): + + _argnames = ('pdf',) + is_Continuous = True + + @property + def set(self): + return self.args[1] + + +def JointRV(symbol, pdf, _set=None): + """ + Create a Joint Random Variable where each of its component is continuous, + given the following: + + Parameters + ========== + + symbol : Symbol + Represents name of the random variable. + pdf : A PDF in terms of indexed symbols of the symbol given + as the first argument + + NOTE + ==== + + As of now, the set for each component for a ``JointRV`` is + equal to the set of all integers, which cannot be changed. + + Examples + ======== + + >>> from sympy import exp, pi, Indexed, S + >>> from sympy.stats import density, JointRV + >>> x1, x2 = (Indexed('x', i) for i in (1, 2)) + >>> pdf = exp(-x1**2/2 + x1 - x2**2/2 - S(1)/2)/(2*pi) + >>> N1 = JointRV('x', pdf) #Multivariate Normal distribution + >>> density(N1)(1, 2) + exp(-2)/(2*pi) + + Returns + ======= + + RandomSymbol + + """ + #TODO: Add support for sets provided by the user + symbol = sympify(symbol) + syms = [i for i in pdf.free_symbols if isinstance(i, Indexed) + and i.base == IndexedBase(symbol)] + syms = tuple(sorted(syms, key = lambda index: index.args[1])) + _set = S.Reals**len(syms) + pdf = Lambda(syms, pdf) + dist = JointDistributionHandmade(pdf, _set) + jrv = JointPSpace(symbol, dist).value + rvs = random_symbols(pdf) + if len(rvs) != 0: + dist = MarginalDistribution(dist, (jrv,)) + return JointPSpace(symbol, dist).value + return jrv + +#------------------------------------------------------------------------------- +# Multivariate Normal distribution --------------------------------------------- + +class MultivariateNormalDistribution(JointDistribution): + _argnames = ('mu', 'sigma') + + is_Continuous=True + + @property + def set(self): + k = self.mu.shape[0] + return S.Reals**k + + @staticmethod + def check(mu, sigma): + _value_check(mu.shape[0] == sigma.shape[0], + "Size of the mean vector and covariance matrix are incorrect.") + #check if covariance matrix is positive semi definite or not. + if not isinstance(sigma, MatrixSymbol): + _value_check(sigma.is_positive_semidefinite, + "The covariance matrix must be positive semi definite. ") + + def pdf(self, *args): + mu, sigma = self.mu, self.sigma + k = mu.shape[0] + if len(args) == 1 and args[0].is_Matrix: + args = args[0] + else: + args = ImmutableMatrix(args) + x = args - mu + density = S.One/sqrt((2*pi)**(k)*det(sigma))*exp( + Rational(-1, 2)*x.transpose()*(sigma.inv()*x)) + return MatrixElement(density, 0, 0) + + def _marginal_distribution(self, indices, sym): + sym = ImmutableMatrix([Indexed(sym, i) for i in indices]) + _mu, _sigma = self.mu, self.sigma + k = self.mu.shape[0] + for i in range(k): + if i not in indices: + _mu = _mu.row_del(i) + _sigma = _sigma.col_del(i) + _sigma = _sigma.row_del(i) + return Lambda(tuple(sym), S.One/sqrt((2*pi)**(len(_mu))*det(_sigma))*exp( + Rational(-1, 2)*(_mu - sym).transpose()*(_sigma.inv()*\ + (_mu - sym)))[0]) + +def MultivariateNormal(name, mu, sigma): + r""" + Creates a continuous random variable with Multivariate Normal + Distribution. + + The density of the multivariate normal distribution can be found at [1]. + + Parameters + ========== + + mu : List representing the mean or the mean vector + sigma : Positive semidefinite square matrix + Represents covariance Matrix. + If `\sigma` is noninvertible then only sampling is supported currently + + Returns + ======= + + RandomSymbol + + Examples + ======== + + >>> from sympy.stats import MultivariateNormal, density, marginal_distribution + >>> from sympy import symbols, MatrixSymbol + >>> X = MultivariateNormal('X', [3, 4], [[2, 1], [1, 2]]) + >>> y, z = symbols('y z') + >>> density(X)(y, z) + sqrt(3)*exp(-y**2/3 + y*z/3 + 2*y/3 - z**2/3 + 5*z/3 - 13/3)/(6*pi) + >>> density(X)(1, 2) + sqrt(3)*exp(-4/3)/(6*pi) + >>> marginal_distribution(X, X[1])(y) + exp(-(y - 4)**2/4)/(2*sqrt(pi)) + >>> marginal_distribution(X, X[0])(y) + exp(-(y - 3)**2/4)/(2*sqrt(pi)) + + The example below shows that it is also possible to use + symbolic parameters to define the MultivariateNormal class. + + >>> n = symbols('n', integer=True, positive=True) + >>> Sg = MatrixSymbol('Sg', n, n) + >>> mu = MatrixSymbol('mu', n, 1) + >>> obs = MatrixSymbol('obs', n, 1) + >>> X = MultivariateNormal('X', mu, Sg) + + The density of a multivariate normal can be + calculated using a matrix argument, as shown below. + + >>> density(X)(obs) + (exp(((1/2)*mu.T - (1/2)*obs.T)*Sg**(-1)*(-mu + obs))/sqrt((2*pi)**n*Determinant(Sg)))[0, 0] + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Multivariate_normal_distribution + + """ + return multivariate_rv(MultivariateNormalDistribution, name, mu, sigma) + +#------------------------------------------------------------------------------- +# Multivariate Laplace distribution -------------------------------------------- + +class MultivariateLaplaceDistribution(JointDistribution): + _argnames = ('mu', 'sigma') + is_Continuous=True + + @property + def set(self): + k = self.mu.shape[0] + return S.Reals**k + + @staticmethod + def check(mu, sigma): + _value_check(mu.shape[0] == sigma.shape[0], + "Size of the mean vector and covariance matrix are incorrect.") + # check if covariance matrix is positive definite or not. + if not isinstance(sigma, MatrixSymbol): + _value_check(sigma.is_positive_definite, + "The covariance matrix must be positive definite. ") + + def pdf(self, *args): + mu, sigma = self.mu, self.sigma + mu_T = mu.transpose() + k = S(mu.shape[0]) + sigma_inv = sigma.inv() + args = ImmutableMatrix(args) + args_T = args.transpose() + x = (mu_T*sigma_inv*mu)[0] + y = (args_T*sigma_inv*args)[0] + v = 1 - k/2 + return (2 * (y/(2 + x))**(v/2) * besselk(v, sqrt((2 + x)*y)) * + exp((args_T * sigma_inv * mu)[0]) / + ((2 * pi)**(k/2) * sqrt(det(sigma)))) + + +def MultivariateLaplace(name, mu, sigma): + """ + Creates a continuous random variable with Multivariate Laplace + Distribution. + + The density of the multivariate Laplace distribution can be found at [1]. + + Parameters + ========== + + mu : List representing the mean or the mean vector + sigma : Positive definite square matrix + Represents covariance Matrix + + Returns + ======= + + RandomSymbol + + Examples + ======== + + >>> from sympy.stats import MultivariateLaplace, density + >>> from sympy import symbols + >>> y, z = symbols('y z') + >>> X = MultivariateLaplace('X', [2, 4], [[3, 1], [1, 3]]) + >>> density(X)(y, z) + sqrt(2)*exp(y/4 + 5*z/4)*besselk(0, sqrt(15*y*(3*y/8 - z/8)/2 + 15*z*(-y/8 + 3*z/8)/2))/(4*pi) + >>> density(X)(1, 2) + sqrt(2)*exp(11/4)*besselk(0, sqrt(165)/4)/(4*pi) + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Multivariate_Laplace_distribution + + """ + return multivariate_rv(MultivariateLaplaceDistribution, name, mu, sigma) + +#------------------------------------------------------------------------------- +# Multivariate StudentT distribution ------------------------------------------- + +class MultivariateTDistribution(JointDistribution): + _argnames = ('mu', 'shape_mat', 'dof') + is_Continuous=True + + @property + def set(self): + k = self.mu.shape[0] + return S.Reals**k + + @staticmethod + def check(mu, sigma, v): + _value_check(mu.shape[0] == sigma.shape[0], + "Size of the location vector and shape matrix are incorrect.") + # check if covariance matrix is positive definite or not. + if not isinstance(sigma, MatrixSymbol): + _value_check(sigma.is_positive_definite, + "The shape matrix must be positive definite. ") + + def pdf(self, *args): + mu, sigma = self.mu, self.shape_mat + v = S(self.dof) + k = S(mu.shape[0]) + sigma_inv = sigma.inv() + args = ImmutableMatrix(args) + x = args - mu + return gamma((k + v)/2)/(gamma(v/2)*(v*pi)**(k/2)*sqrt(det(sigma)))\ + *(1 + 1/v*(x.transpose()*sigma_inv*x)[0])**((-v - k)/2) + +def MultivariateT(syms, mu, sigma, v): + """ + Creates a joint random variable with multivariate T-distribution. + + Parameters + ========== + + syms : A symbol/str + For identifying the random variable. + mu : A list/matrix + Representing the location vector + sigma : The shape matrix for the distribution + + Examples + ======== + + >>> from sympy.stats import density, MultivariateT + >>> from sympy import Symbol + + >>> x = Symbol("x") + >>> X = MultivariateT("x", [1, 1], [[1, 0], [0, 1]], 2) + + >>> density(X)(1, 2) + 2/(9*pi) + + Returns + ======= + + RandomSymbol + + """ + return multivariate_rv(MultivariateTDistribution, syms, mu, sigma, v) + + +#------------------------------------------------------------------------------- +# Multivariate Normal Gamma distribution --------------------------------------- + +class NormalGammaDistribution(JointDistribution): + + _argnames = ('mu', 'lamda', 'alpha', 'beta') + is_Continuous=True + + @staticmethod + def check(mu, lamda, alpha, beta): + _value_check(mu.is_real, "Location must be real.") + _value_check(lamda > 0, "Lambda must be positive") + _value_check(alpha > 0, "alpha must be positive") + _value_check(beta > 0, "beta must be positive") + + @property + def set(self): + return S.Reals*Interval(0, S.Infinity) + + def pdf(self, x, tau): + beta, alpha, lamda = self.beta, self.alpha, self.lamda + mu = self.mu + + return beta**alpha*sqrt(lamda)/(gamma(alpha)*sqrt(2*pi))*\ + tau**(alpha - S.Half)*exp(-1*beta*tau)*\ + exp(-1*(lamda*tau*(x - mu)**2)/S(2)) + + def _marginal_distribution(self, indices, *sym): + if len(indices) == 2: + return self.pdf(*sym) + if indices[0] == 0: + #For marginal over `x`, return non-standardized Student-T's + #distribution + x = sym[0] + v, mu, sigma = self.alpha - S.Half, self.mu, \ + S(self.beta)/(self.lamda * self.alpha) + return Lambda(sym, gamma((v + 1)/2)/(gamma(v/2)*sqrt(pi*v)*sigma)*\ + (1 + 1/v*((x - mu)/sigma)**2)**((-v -1)/2)) + #For marginal over `tau`, return Gamma distribution as per construction + from sympy.stats.crv_types import GammaDistribution + return Lambda(sym, GammaDistribution(self.alpha, self.beta)(sym[0])) + +def NormalGamma(sym, mu, lamda, alpha, beta): + """ + Creates a bivariate joint random variable with multivariate Normal gamma + distribution. + + Parameters + ========== + + sym : A symbol/str + For identifying the random variable. + mu : A real number + The mean of the normal distribution + lamda : A positive integer + Parameter of joint distribution + alpha : A positive integer + Parameter of joint distribution + beta : A positive integer + Parameter of joint distribution + + Returns + ======= + + RandomSymbol + + Examples + ======== + + >>> from sympy.stats import density, NormalGamma + >>> from sympy import symbols + + >>> X = NormalGamma('x', 0, 1, 2, 3) + >>> y, z = symbols('y z') + + >>> density(X)(y, z) + 9*sqrt(2)*z**(3/2)*exp(-3*z)*exp(-y**2*z/2)/(2*sqrt(pi)) + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Normal-gamma_distribution + + """ + return multivariate_rv(NormalGammaDistribution, sym, mu, lamda, alpha, beta) + +#------------------------------------------------------------------------------- +# Multivariate Beta/Dirichlet distribution ------------------------------------- + +class MultivariateBetaDistribution(JointDistribution): + + _argnames = ('alpha',) + is_Continuous = True + + @staticmethod + def check(alpha): + _value_check(len(alpha) >= 2, "At least two categories should be passed.") + for a_k in alpha: + _value_check((a_k > 0) != False, "Each concentration parameter" + " should be positive.") + + @property + def set(self): + k = len(self.alpha) + return Interval(0, 1)**k + + def pdf(self, *syms): + alpha = self.alpha + B = Mul.fromiter(map(gamma, alpha))/gamma(Add(*alpha)) + return Mul.fromiter(sym**(a_k - 1) for a_k, sym in zip(alpha, syms))/B + +def MultivariateBeta(syms, *alpha): + """ + Creates a continuous random variable with Dirichlet/Multivariate Beta + Distribution. + + The density of the Dirichlet distribution can be found at [1]. + + Parameters + ========== + + alpha : Positive real numbers + Signifies concentration numbers. + + Returns + ======= + + RandomSymbol + + Examples + ======== + + >>> from sympy.stats import density, MultivariateBeta, marginal_distribution + >>> from sympy import Symbol + >>> a1 = Symbol('a1', positive=True) + >>> a2 = Symbol('a2', positive=True) + >>> B = MultivariateBeta('B', [a1, a2]) + >>> C = MultivariateBeta('C', a1, a2) + >>> x = Symbol('x') + >>> y = Symbol('y') + >>> density(B)(x, y) + x**(a1 - 1)*y**(a2 - 1)*gamma(a1 + a2)/(gamma(a1)*gamma(a2)) + >>> marginal_distribution(C, C[0])(x) + x**(a1 - 1)*gamma(a1 + a2)/(a2*gamma(a1)*gamma(a2)) + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Dirichlet_distribution + .. [2] https://mathworld.wolfram.com/DirichletDistribution.html + + """ + if not isinstance(alpha[0], list): + alpha = (list(alpha),) + return multivariate_rv(MultivariateBetaDistribution, syms, alpha[0]) + +Dirichlet = MultivariateBeta + +#------------------------------------------------------------------------------- +# Multivariate Ewens distribution ---------------------------------------------- + +class MultivariateEwensDistribution(JointDistribution): + + _argnames = ('n', 'theta') + is_Discrete = True + is_Continuous = False + + @staticmethod + def check(n, theta): + _value_check((n > 0), + "sample size should be positive integer.") + _value_check(theta.is_positive, "mutation rate should be positive.") + + @property + def set(self): + if not isinstance(self.n, Integer): + i = Symbol('i', integer=True, positive=True) + return Product(Intersection(S.Naturals0, Interval(0, self.n//i)), + (i, 1, self.n)) + prod_set = Range(0, self.n + 1) + for i in range(2, self.n + 1): + prod_set *= Range(0, self.n//i + 1) + return prod_set.flatten() + + def pdf(self, *syms): + n, theta = self.n, self.theta + condi = isinstance(self.n, Integer) + if not (isinstance(syms[0], IndexedBase) or condi): + raise ValueError("Please use IndexedBase object for syms as " + "the dimension is symbolic") + term_1 = factorial(n)/rf(theta, n) + if condi: + term_2 = Mul.fromiter(theta**syms[j]/((j+1)**syms[j]*factorial(syms[j])) + for j in range(n)) + cond = Eq(sum((k + 1)*syms[k] for k in range(n)), n) + return Piecewise((term_1 * term_2, cond), (0, True)) + syms = syms[0] + j, k = symbols('j, k', positive=True, integer=True) + term_2 = Product(theta**syms[j]/((j+1)**syms[j]*factorial(syms[j])), + (j, 0, n - 1)) + cond = Eq(Sum((k + 1)*syms[k], (k, 0, n - 1)), n) + return Piecewise((term_1 * term_2, cond), (0, True)) + + +def MultivariateEwens(syms, n, theta): + """ + Creates a discrete random variable with Multivariate Ewens + Distribution. + + The density of the said distribution can be found at [1]. + + Parameters + ========== + + n : Positive integer + Size of the sample or the integer whose partitions are considered + theta : Positive real number + Denotes Mutation rate + + Returns + ======= + + RandomSymbol + + Examples + ======== + + >>> from sympy.stats import density, marginal_distribution, MultivariateEwens + >>> from sympy import Symbol + >>> a1 = Symbol('a1', positive=True) + >>> a2 = Symbol('a2', positive=True) + >>> ed = MultivariateEwens('E', 2, 1) + >>> density(ed)(a1, a2) + Piecewise((1/(2**a2*factorial(a1)*factorial(a2)), Eq(a1 + 2*a2, 2)), (0, True)) + >>> marginal_distribution(ed, ed[0])(a1) + Piecewise((1/factorial(a1), Eq(a1, 2)), (0, True)) + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Ewens%27s_sampling_formula + .. [2] https://www.jstor.org/stable/24780825 + """ + return multivariate_rv(MultivariateEwensDistribution, syms, n, theta) + +#------------------------------------------------------------------------------- +# Generalized Multivariate Log Gamma distribution ------------------------------ + +class GeneralizedMultivariateLogGammaDistribution(JointDistribution): + + _argnames = ('delta', 'v', 'lamda', 'mu') + is_Continuous=True + + def check(self, delta, v, l, mu): + _value_check((delta >= 0, delta <= 1), "delta must be in range [0, 1].") + _value_check((v > 0), "v must be positive") + for lk in l: + _value_check((lk > 0), "lamda must be a positive vector.") + for muk in mu: + _value_check((muk > 0), "mu must be a positive vector.") + _value_check(len(l) > 1,"the distribution should have at least" + " two random variables.") + + @property + def set(self): + return S.Reals**len(self.lamda) + + def pdf(self, *y): + d, v, l, mu = self.delta, self.v, self.lamda, self.mu + n = Symbol('n', negative=False, integer=True) + k = len(l) + sterm1 = Pow((1 - d), n)/\ + ((gamma(v + n)**(k - 1))*gamma(v)*gamma(n + 1)) + sterm2 = Mul.fromiter(mui*li**(-v - n) for mui, li in zip(mu, l)) + term1 = sterm1 * sterm2 + sterm3 = (v + n) * sum(mui * yi for mui, yi in zip(mu, y)) + sterm4 = sum(exp(mui * yi)/li for (mui, yi, li) in zip(mu, y, l)) + term2 = exp(sterm3 - sterm4) + return Pow(d, v) * Sum(term1 * term2, (n, 0, S.Infinity)) + +def GeneralizedMultivariateLogGamma(syms, delta, v, lamda, mu): + """ + Creates a joint random variable with generalized multivariate log gamma + distribution. + + The joint pdf can be found at [1]. + + Parameters + ========== + + syms : list/tuple/set of symbols for identifying each component + delta : A constant in range $[0, 1]$ + v : Positive real number + lamda : List of positive real numbers + mu : List of positive real numbers + + Returns + ======= + + RandomSymbol + + Examples + ======== + + >>> from sympy.stats import density + >>> from sympy.stats.joint_rv_types import GeneralizedMultivariateLogGamma + >>> from sympy import symbols, S + >>> v = 1 + >>> l, mu = [1, 1, 1], [1, 1, 1] + >>> d = S.Half + >>> y = symbols('y_1:4', positive=True) + >>> Gd = GeneralizedMultivariateLogGamma('G', d, v, l, mu) + >>> density(Gd)(y[0], y[1], y[2]) + Sum(exp((n + 1)*(y_1 + y_2 + y_3) - exp(y_1) - exp(y_2) - + exp(y_3))/(2**n*gamma(n + 1)**3), (n, 0, oo))/2 + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Generalized_multivariate_log-gamma_distribution + .. [2] https://www.researchgate.net/publication/234137346_On_a_multivariate_log-gamma_distribution_and_the_use_of_the_distribution_in_the_Bayesian_analysis + + Note + ==== + + If the GeneralizedMultivariateLogGamma is too long to type use, + + >>> from sympy.stats.joint_rv_types import GeneralizedMultivariateLogGamma as GMVLG + >>> Gd = GMVLG('G', d, v, l, mu) + + If you want to pass the matrix omega instead of the constant delta, then use + ``GeneralizedMultivariateLogGammaOmega``. + + """ + return multivariate_rv(GeneralizedMultivariateLogGammaDistribution, + syms, delta, v, lamda, mu) + +def GeneralizedMultivariateLogGammaOmega(syms, omega, v, lamda, mu): + """ + Extends GeneralizedMultivariateLogGamma. + + Parameters + ========== + + syms : list/tuple/set of symbols + For identifying each component + omega : A square matrix + Every element of square matrix must be absolute value of + square root of correlation coefficient + v : Positive real number + lamda : List of positive real numbers + mu : List of positive real numbers + + Returns + ======= + + RandomSymbol + + Examples + ======== + + >>> from sympy.stats import density + >>> from sympy.stats.joint_rv_types import GeneralizedMultivariateLogGammaOmega + >>> from sympy import Matrix, symbols, S + >>> omega = Matrix([[1, S.Half, S.Half], [S.Half, 1, S.Half], [S.Half, S.Half, 1]]) + >>> v = 1 + >>> l, mu = [1, 1, 1], [1, 1, 1] + >>> G = GeneralizedMultivariateLogGammaOmega('G', omega, v, l, mu) + >>> y = symbols('y_1:4', positive=True) + >>> density(G)(y[0], y[1], y[2]) + sqrt(2)*Sum((1 - sqrt(2)/2)**n*exp((n + 1)*(y_1 + y_2 + y_3) - exp(y_1) - + exp(y_2) - exp(y_3))/gamma(n + 1)**3, (n, 0, oo))/2 + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Generalized_multivariate_log-gamma_distribution + .. [2] https://www.researchgate.net/publication/234137346_On_a_multivariate_log-gamma_distribution_and_the_use_of_the_distribution_in_the_Bayesian_analysis + + Notes + ===== + + If the GeneralizedMultivariateLogGammaOmega is too long to type use, + + >>> from sympy.stats.joint_rv_types import GeneralizedMultivariateLogGammaOmega as GMVLGO + >>> G = GMVLGO('G', omega, v, l, mu) + + """ + _value_check((omega.is_square, isinstance(omega, Matrix)), "omega must be a" + " square matrix") + for val in omega.values(): + _value_check((val >= 0, val <= 1), + "all values in matrix must be between 0 and 1(both inclusive).") + _value_check(omega.diagonal().equals(ones(1, omega.shape[0])), + "all the elements of diagonal should be 1.") + _value_check((omega.shape[0] == len(lamda), len(lamda) == len(mu)), + "lamda, mu should be of same length and omega should " + " be of shape (length of lamda, length of mu)") + _value_check(len(lamda) > 1,"the distribution should have at least" + " two random variables.") + delta = Pow(Rational(omega.det()), Rational(1, len(lamda) - 1)) + return GeneralizedMultivariateLogGamma(syms, delta, v, lamda, mu) + + +#------------------------------------------------------------------------------- +# Multinomial distribution ----------------------------------------------------- + +class MultinomialDistribution(JointDistribution): + + _argnames = ('n', 'p') + is_Continuous=False + is_Discrete = True + + @staticmethod + def check(n, p): + _value_check(n > 0, + "number of trials must be a positive integer") + for p_k in p: + _value_check((p_k >= 0, p_k <= 1), + "probability must be in range [0, 1]") + _value_check(Eq(sum(p), 1), + "probabilities must sum to 1") + + @property + def set(self): + return Intersection(S.Naturals0, Interval(0, self.n))**len(self.p) + + def pdf(self, *x): + n, p = self.n, self.p + term_1 = factorial(n)/Mul.fromiter(factorial(x_k) for x_k in x) + term_2 = Mul.fromiter(p_k**x_k for p_k, x_k in zip(p, x)) + return Piecewise((term_1 * term_2, Eq(sum(x), n)), (0, True)) + +def Multinomial(syms, n, *p): + """ + Creates a discrete random variable with Multinomial Distribution. + + The density of the said distribution can be found at [1]. + + Parameters + ========== + + n : Positive integer + Represents number of trials + p : List of event probabilities + Must be in the range of $[0, 1]$. + + Returns + ======= + + RandomSymbol + + Examples + ======== + + >>> from sympy.stats import density, Multinomial, marginal_distribution + >>> from sympy import symbols + >>> x1, x2, x3 = symbols('x1, x2, x3', nonnegative=True, integer=True) + >>> p1, p2, p3 = symbols('p1, p2, p3', positive=True) + >>> M = Multinomial('M', 3, p1, p2, p3) + >>> density(M)(x1, x2, x3) + Piecewise((6*p1**x1*p2**x2*p3**x3/(factorial(x1)*factorial(x2)*factorial(x3)), + Eq(x1 + x2 + x3, 3)), (0, True)) + >>> marginal_distribution(M, M[0])(x1).subs(x1, 1) + 3*p1*p2**2 + 6*p1*p2*p3 + 3*p1*p3**2 + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Multinomial_distribution + .. [2] https://mathworld.wolfram.com/MultinomialDistribution.html + + """ + if not isinstance(p[0], list): + p = (list(p), ) + return multivariate_rv(MultinomialDistribution, syms, n, p[0]) + +#------------------------------------------------------------------------------- +# Negative Multinomial Distribution -------------------------------------------- + +class NegativeMultinomialDistribution(JointDistribution): + + _argnames = ('k0', 'p') + is_Continuous=False + is_Discrete = True + + @staticmethod + def check(k0, p): + _value_check(k0 > 0, + "number of failures must be a positive integer") + for p_k in p: + _value_check((p_k >= 0, p_k <= 1), + "probability must be in range [0, 1].") + _value_check(sum(p) <= 1, + "success probabilities must not be greater than 1.") + + @property + def set(self): + return Range(0, S.Infinity)**len(self.p) + + def pdf(self, *k): + k0, p = self.k0, self.p + term_1 = (gamma(k0 + sum(k))*(1 - sum(p))**k0)/gamma(k0) + term_2 = Mul.fromiter(pi**ki/factorial(ki) for pi, ki in zip(p, k)) + return term_1 * term_2 + +def NegativeMultinomial(syms, k0, *p): + """ + Creates a discrete random variable with Negative Multinomial Distribution. + + The density of the said distribution can be found at [1]. + + Parameters + ========== + + k0 : positive integer + Represents number of failures before the experiment is stopped + p : List of event probabilities + Must be in the range of $[0, 1]$ + + Returns + ======= + + RandomSymbol + + Examples + ======== + + >>> from sympy.stats import density, NegativeMultinomial, marginal_distribution + >>> from sympy import symbols + >>> x1, x2, x3 = symbols('x1, x2, x3', nonnegative=True, integer=True) + >>> p1, p2, p3 = symbols('p1, p2, p3', positive=True) + >>> N = NegativeMultinomial('M', 3, p1, p2, p3) + >>> N_c = NegativeMultinomial('M', 3, 0.1, 0.1, 0.1) + >>> density(N)(x1, x2, x3) + p1**x1*p2**x2*p3**x3*(-p1 - p2 - p3 + 1)**3*gamma(x1 + x2 + + x3 + 3)/(2*factorial(x1)*factorial(x2)*factorial(x3)) + >>> marginal_distribution(N_c, N_c[0])(1).evalf().round(2) + 0.25 + + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Negative_multinomial_distribution + .. [2] https://mathworld.wolfram.com/NegativeBinomialDistribution.html + + """ + if not isinstance(p[0], list): + p = (list(p), ) + return multivariate_rv(NegativeMultinomialDistribution, syms, k0, p[0]) diff --git a/lib/python3.12/site-packages/sympy/stats/matrix_distributions.py b/lib/python3.12/site-packages/sympy/stats/matrix_distributions.py new file mode 100644 index 0000000000000000000000000000000000000000..9a43c0226bc25702211a910ebbe30e280ad0cf50 --- /dev/null +++ b/lib/python3.12/site-packages/sympy/stats/matrix_distributions.py @@ -0,0 +1,610 @@ +from math import prod + +from sympy.core.basic import Basic +from sympy.core.numbers import pi +from sympy.core.singleton import S +from sympy.functions.elementary.exponential import exp +from sympy.functions.special.gamma_functions import multigamma +from sympy.core.sympify import sympify, _sympify +from sympy.matrices import (ImmutableMatrix, Inverse, Trace, Determinant, + MatrixSymbol, MatrixBase, Transpose, MatrixSet, + matrix2numpy) +from sympy.stats.rv import (_value_check, RandomMatrixSymbol, NamedArgsMixin, PSpace, + _symbol_converter, MatrixDomain, Distribution) +from sympy.external import import_module + + +################################################################################ +#------------------------Matrix Probability Space------------------------------# +################################################################################ +class MatrixPSpace(PSpace): + """ + Represents probability space for + Matrix Distributions. + """ + def __new__(cls, sym, distribution, dim_n, dim_m): + sym = _symbol_converter(sym) + dim_n, dim_m = _sympify(dim_n), _sympify(dim_m) + if not (dim_n.is_integer and dim_m.is_integer): + raise ValueError("Dimensions should be integers") + return Basic.__new__(cls, sym, distribution, dim_n, dim_m) + + distribution = property(lambda self: self.args[1]) + symbol = property(lambda self: self.args[0]) + + @property + def domain(self): + return MatrixDomain(self.symbol, self.distribution.set) + + @property + def value(self): + return RandomMatrixSymbol(self.symbol, self.args[2], self.args[3], self) + + @property + def values(self): + return {self.value} + + def compute_density(self, expr, *args): + rms = expr.atoms(RandomMatrixSymbol) + if len(rms) > 1 or (not isinstance(expr, RandomMatrixSymbol)): + raise NotImplementedError("Currently, no algorithm has been " + "implemented to handle general expressions containing " + "multiple matrix distributions.") + return self.distribution.pdf(expr) + + def sample(self, size=(), library='scipy', seed=None): + """ + Internal sample method + + Returns dictionary mapping RandomMatrixSymbol to realization value. + """ + return {self.value: self.distribution.sample(size, library=library, seed=seed)} + + +def rv(symbol, cls, args): + args = list(map(sympify, args)) + dist = cls(*args) + dist.check(*args) + dim = dist.dimension + pspace = MatrixPSpace(symbol, dist, dim[0], dim[1]) + return pspace.value + + +class SampleMatrixScipy: + """Returns the sample from scipy of the given distribution""" + def __new__(cls, dist, size, seed=None): + return cls._sample_scipy(dist, size, seed) + + @classmethod + def _sample_scipy(cls, dist, size, seed): + """Sample from SciPy.""" + + from scipy import stats as scipy_stats + import numpy + scipy_rv_map = { + 'WishartDistribution': lambda dist, size, rand_state: scipy_stats.wishart.rvs( + df=int(dist.n), scale=matrix2numpy(dist.scale_matrix, float), size=size), + 'MatrixNormalDistribution': lambda dist, size, rand_state: scipy_stats.matrix_normal.rvs( + mean=matrix2numpy(dist.location_matrix, float), + rowcov=matrix2numpy(dist.scale_matrix_1, float), + colcov=matrix2numpy(dist.scale_matrix_2, float), size=size, random_state=rand_state) + } + + sample_shape = { + 'WishartDistribution': lambda dist: dist.scale_matrix.shape, + 'MatrixNormalDistribution' : lambda dist: dist.location_matrix.shape + } + + dist_list = scipy_rv_map.keys() + + if dist.__class__.__name__ not in dist_list: + return None + + if seed is None or isinstance(seed, int): + rand_state = numpy.random.default_rng(seed=seed) + else: + rand_state = seed + samp = scipy_rv_map[dist.__class__.__name__](dist, prod(size), rand_state) + return samp.reshape(size + sample_shape[dist.__class__.__name__](dist)) + + +class SampleMatrixNumpy: + """Returns the sample from numpy of the given distribution""" + + ### TODO: Add tests after adding matrix distributions in numpy_rv_map + def __new__(cls, dist, size, seed=None): + return cls._sample_numpy(dist, size, seed) + + @classmethod + def _sample_numpy(cls, dist, size, seed): + """Sample from NumPy.""" + + numpy_rv_map = { + } + + sample_shape = { + } + + dist_list = numpy_rv_map.keys() + + if dist.__class__.__name__ not in dist_list: + return None + + import numpy + if seed is None or isinstance(seed, int): + rand_state = numpy.random.default_rng(seed=seed) + else: + rand_state = seed + samp = numpy_rv_map[dist.__class__.__name__](dist, prod(size), rand_state) + return samp.reshape(size + sample_shape[dist.__class__.__name__](dist)) + + +class SampleMatrixPymc: + """Returns the sample from pymc of the given distribution""" + + def __new__(cls, dist, size, seed=None): + return cls._sample_pymc(dist, size, seed) + + @classmethod + def _sample_pymc(cls, dist, size, seed): + """Sample from PyMC.""" + + try: + import pymc + except ImportError: + import pymc3 as pymc + pymc_rv_map = { + 'MatrixNormalDistribution': lambda dist: pymc.MatrixNormal('X', + mu=matrix2numpy(dist.location_matrix, float), + rowcov=matrix2numpy(dist.scale_matrix_1, float), + colcov=matrix2numpy(dist.scale_matrix_2, float), + shape=dist.location_matrix.shape), + 'WishartDistribution': lambda dist: pymc.WishartBartlett('X', + nu=int(dist.n), S=matrix2numpy(dist.scale_matrix, float)) + } + + sample_shape = { + 'WishartDistribution': lambda dist: dist.scale_matrix.shape, + 'MatrixNormalDistribution' : lambda dist: dist.location_matrix.shape + } + + dist_list = pymc_rv_map.keys() + + if dist.__class__.__name__ not in dist_list: + return None + import logging + logging.getLogger("pymc").setLevel(logging.ERROR) + with pymc.Model(): + pymc_rv_map[dist.__class__.__name__](dist) + samps = pymc.sample(draws=prod(size), chains=1, progressbar=False, random_seed=seed, return_inferencedata=False, compute_convergence_checks=False)['X'] + return samps.reshape(size + sample_shape[dist.__class__.__name__](dist)) + +_get_sample_class_matrixrv = { + 'scipy': SampleMatrixScipy, + 'pymc3': SampleMatrixPymc, + 'pymc': SampleMatrixPymc, + 'numpy': SampleMatrixNumpy +} + +################################################################################ +#-------------------------Matrix Distribution----------------------------------# +################################################################################ + +class MatrixDistribution(Distribution, NamedArgsMixin): + """ + Abstract class for Matrix Distribution. + """ + def __new__(cls, *args): + args = [ImmutableMatrix(arg) if isinstance(arg, list) + else _sympify(arg) for arg in args] + return Basic.__new__(cls, *args) + + @staticmethod + def check(*args): + pass + + def __call__(self, expr): + if isinstance(expr, list): + expr = ImmutableMatrix(expr) + return self.pdf(expr) + + def sample(self, size=(), library='scipy', seed=None): + """ + Internal sample method + + Returns dictionary mapping RandomSymbol to realization value. + """ + + libraries = ['scipy', 'numpy', 'pymc3', 'pymc'] + if library not in libraries: + raise NotImplementedError("Sampling from %s is not supported yet." + % str(library)) + if not import_module(library): + raise ValueError("Failed to import %s" % library) + + samps = _get_sample_class_matrixrv[library](self, size, seed) + + if samps is not None: + return samps + raise NotImplementedError( + "Sampling for %s is not currently implemented from %s" + % (self.__class__.__name__, library) + ) + +################################################################################ +#------------------------Matrix Distribution Types-----------------------------# +################################################################################ + +#------------------------------------------------------------------------------- +# Matrix Gamma distribution ---------------------------------------------------- + +class MatrixGammaDistribution(MatrixDistribution): + + _argnames = ('alpha', 'beta', 'scale_matrix') + + @staticmethod + def check(alpha, beta, scale_matrix): + if not isinstance(scale_matrix, MatrixSymbol): + _value_check(scale_matrix.is_positive_definite, "The shape " + "matrix must be positive definite.") + _value_check(scale_matrix.is_square, "Should " + "be square matrix") + _value_check(alpha.is_positive, "Shape parameter should be positive.") + _value_check(beta.is_positive, "Scale parameter should be positive.") + + @property + def set(self): + k = self.scale_matrix.shape[0] + return MatrixSet(k, k, S.Reals) + + @property + def dimension(self): + return self.scale_matrix.shape + + def pdf(self, x): + alpha, beta, scale_matrix = self.alpha, self.beta, self.scale_matrix + p = scale_matrix.shape[0] + if isinstance(x, list): + x = ImmutableMatrix(x) + if not isinstance(x, (MatrixBase, MatrixSymbol)): + raise ValueError("%s should be an isinstance of Matrix " + "or MatrixSymbol" % str(x)) + sigma_inv_x = - Inverse(scale_matrix)*x / beta + term1 = exp(Trace(sigma_inv_x))/((beta**(p*alpha)) * multigamma(alpha, p)) + term2 = (Determinant(scale_matrix))**(-alpha) + term3 = (Determinant(x))**(alpha - S(p + 1)/2) + return term1 * term2 * term3 + +def MatrixGamma(symbol, alpha, beta, scale_matrix): + """ + Creates a random variable with Matrix Gamma Distribution. + + The density of the said distribution can be found at [1]. + + Parameters + ========== + + alpha: Positive Real number + Shape Parameter + beta: Positive Real number + Scale Parameter + scale_matrix: Positive definite real square matrix + Scale Matrix + + Returns + ======= + + RandomSymbol + + Examples + ======== + + >>> from sympy.stats import density, MatrixGamma + >>> from sympy import MatrixSymbol, symbols + >>> a, b = symbols('a b', positive=True) + >>> M = MatrixGamma('M', a, b, [[2, 1], [1, 2]]) + >>> X = MatrixSymbol('X', 2, 2) + >>> density(M)(X).doit() + exp(Trace(Matrix([ + [-2/3, 1/3], + [ 1/3, -2/3]])*X)/b)*Determinant(X)**(a - 3/2)/(3**a*sqrt(pi)*b**(2*a)*gamma(a)*gamma(a - 1/2)) + >>> density(M)([[1, 0], [0, 1]]).doit() + exp(-4/(3*b))/(3**a*sqrt(pi)*b**(2*a)*gamma(a)*gamma(a - 1/2)) + + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Matrix_gamma_distribution + + """ + if isinstance(scale_matrix, list): + scale_matrix = ImmutableMatrix(scale_matrix) + return rv(symbol, MatrixGammaDistribution, (alpha, beta, scale_matrix)) + +#------------------------------------------------------------------------------- +# Wishart Distribution --------------------------------------------------------- + +class WishartDistribution(MatrixDistribution): + + _argnames = ('n', 'scale_matrix') + + @staticmethod + def check(n, scale_matrix): + if not isinstance(scale_matrix, MatrixSymbol): + _value_check(scale_matrix.is_positive_definite, "The shape " + "matrix must be positive definite.") + _value_check(scale_matrix.is_square, "Should " + "be square matrix") + _value_check(n.is_positive, "Shape parameter should be positive.") + + @property + def set(self): + k = self.scale_matrix.shape[0] + return MatrixSet(k, k, S.Reals) + + @property + def dimension(self): + return self.scale_matrix.shape + + def pdf(self, x): + n, scale_matrix = self.n, self.scale_matrix + p = scale_matrix.shape[0] + if isinstance(x, list): + x = ImmutableMatrix(x) + if not isinstance(x, (MatrixBase, MatrixSymbol)): + raise ValueError("%s should be an isinstance of Matrix " + "or MatrixSymbol" % str(x)) + sigma_inv_x = - Inverse(scale_matrix)*x / S(2) + term1 = exp(Trace(sigma_inv_x))/((2**(p*n/S(2))) * multigamma(n/S(2), p)) + term2 = (Determinant(scale_matrix))**(-n/S(2)) + term3 = (Determinant(x))**(S(n - p - 1)/2) + return term1 * term2 * term3 + +def Wishart(symbol, n, scale_matrix): + """ + Creates a random variable with Wishart Distribution. + + The density of the said distribution can be found at [1]. + + Parameters + ========== + + n: Positive Real number + Represents degrees of freedom + scale_matrix: Positive definite real square matrix + Scale Matrix + + Returns + ======= + + RandomSymbol + + Examples + ======== + + >>> from sympy.stats import density, Wishart + >>> from sympy import MatrixSymbol, symbols + >>> n = symbols('n', positive=True) + >>> W = Wishart('W', n, [[2, 1], [1, 2]]) + >>> X = MatrixSymbol('X', 2, 2) + >>> density(W)(X).doit() + exp(Trace(Matrix([ + [-1/3, 1/6], + [ 1/6, -1/3]])*X))*Determinant(X)**(n/2 - 3/2)/(2**n*3**(n/2)*sqrt(pi)*gamma(n/2)*gamma(n/2 - 1/2)) + >>> density(W)([[1, 0], [0, 1]]).doit() + exp(-2/3)/(2**n*3**(n/2)*sqrt(pi)*gamma(n/2)*gamma(n/2 - 1/2)) + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Wishart_distribution + + """ + if isinstance(scale_matrix, list): + scale_matrix = ImmutableMatrix(scale_matrix) + return rv(symbol, WishartDistribution, (n, scale_matrix)) + +#------------------------------------------------------------------------------- +# Matrix Normal distribution --------------------------------------------------- + +class MatrixNormalDistribution(MatrixDistribution): + + _argnames = ('location_matrix', 'scale_matrix_1', 'scale_matrix_2') + + @staticmethod + def check(location_matrix, scale_matrix_1, scale_matrix_2): + if not isinstance(scale_matrix_1, MatrixSymbol): + _value_check(scale_matrix_1.is_positive_definite, "The shape " + "matrix must be positive definite.") + if not isinstance(scale_matrix_2, MatrixSymbol): + _value_check(scale_matrix_2.is_positive_definite, "The shape " + "matrix must be positive definite.") + _value_check(scale_matrix_1.is_square, "Scale matrix 1 should be " + "be square matrix") + _value_check(scale_matrix_2.is_square, "Scale matrix 2 should be " + "be square matrix") + n = location_matrix.shape[0] + p = location_matrix.shape[1] + _value_check(scale_matrix_1.shape[0] == n, "Scale matrix 1 should be" + " of shape %s x %s"% (str(n), str(n))) + _value_check(scale_matrix_2.shape[0] == p, "Scale matrix 2 should be" + " of shape %s x %s"% (str(p), str(p))) + + @property + def set(self): + n, p = self.location_matrix.shape + return MatrixSet(n, p, S.Reals) + + @property + def dimension(self): + return self.location_matrix.shape + + def pdf(self, x): + M, U, V = self.location_matrix, self.scale_matrix_1, self.scale_matrix_2 + n, p = M.shape + if isinstance(x, list): + x = ImmutableMatrix(x) + if not isinstance(x, (MatrixBase, MatrixSymbol)): + raise ValueError("%s should be an isinstance of Matrix " + "or MatrixSymbol" % str(x)) + term1 = Inverse(V)*Transpose(x - M)*Inverse(U)*(x - M) + num = exp(-Trace(term1)/S(2)) + den = (2*pi)**(S(n*p)/2) * Determinant(U)**(S(p)/2) * Determinant(V)**(S(n)/2) + return num/den + +def MatrixNormal(symbol, location_matrix, scale_matrix_1, scale_matrix_2): + """ + Creates a random variable with Matrix Normal Distribution. + + The density of the said distribution can be found at [1]. + + Parameters + ========== + + location_matrix: Real ``n x p`` matrix + Represents degrees of freedom + scale_matrix_1: Positive definite matrix + Scale Matrix of shape ``n x n`` + scale_matrix_2: Positive definite matrix + Scale Matrix of shape ``p x p`` + + Returns + ======= + + RandomSymbol + + Examples + ======== + + >>> from sympy import MatrixSymbol + >>> from sympy.stats import density, MatrixNormal + >>> M = MatrixNormal('M', [[1, 2]], [1], [[1, 0], [0, 1]]) + >>> X = MatrixSymbol('X', 1, 2) + >>> density(M)(X).doit() + exp(-Trace((Matrix([ + [-1], + [-2]]) + X.T)*(Matrix([[-1, -2]]) + X))/2)/(2*pi) + >>> density(M)([[3, 4]]).doit() + exp(-4)/(2*pi) + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Matrix_normal_distribution + + """ + if isinstance(location_matrix, list): + location_matrix = ImmutableMatrix(location_matrix) + if isinstance(scale_matrix_1, list): + scale_matrix_1 = ImmutableMatrix(scale_matrix_1) + if isinstance(scale_matrix_2, list): + scale_matrix_2 = ImmutableMatrix(scale_matrix_2) + args = (location_matrix, scale_matrix_1, scale_matrix_2) + return rv(symbol, MatrixNormalDistribution, args) + +#------------------------------------------------------------------------------- +# Matrix Student's T distribution --------------------------------------------------- + +class MatrixStudentTDistribution(MatrixDistribution): + + _argnames = ('nu', 'location_matrix', 'scale_matrix_1', 'scale_matrix_2') + + @staticmethod + def check(nu, location_matrix, scale_matrix_1, scale_matrix_2): + if not isinstance(scale_matrix_1, MatrixSymbol): + _value_check(scale_matrix_1.is_positive_definite != False, "The shape " + "matrix must be positive definite.") + if not isinstance(scale_matrix_2, MatrixSymbol): + _value_check(scale_matrix_2.is_positive_definite != False, "The shape " + "matrix must be positive definite.") + _value_check(scale_matrix_1.is_square != False, "Scale matrix 1 should be " + "be square matrix") + _value_check(scale_matrix_2.is_square != False, "Scale matrix 2 should be " + "be square matrix") + n = location_matrix.shape[0] + p = location_matrix.shape[1] + _value_check(scale_matrix_1.shape[0] == p, "Scale matrix 1 should be" + " of shape %s x %s" % (str(p), str(p))) + _value_check(scale_matrix_2.shape[0] == n, "Scale matrix 2 should be" + " of shape %s x %s" % (str(n), str(n))) + _value_check(nu.is_positive != False, "Degrees of freedom must be positive") + + @property + def set(self): + n, p = self.location_matrix.shape + return MatrixSet(n, p, S.Reals) + + @property + def dimension(self): + return self.location_matrix.shape + + def pdf(self, x): + from sympy.matrices.dense import eye + if isinstance(x, list): + x = ImmutableMatrix(x) + if not isinstance(x, (MatrixBase, MatrixSymbol)): + raise ValueError("%s should be an isinstance of Matrix " + "or MatrixSymbol" % str(x)) + nu, M, Omega, Sigma = self.nu, self.location_matrix, self.scale_matrix_1, self.scale_matrix_2 + n, p = M.shape + + K = multigamma((nu + n + p - 1)/2, p) * Determinant(Omega)**(-n/2) * Determinant(Sigma)**(-p/2) \ + / ((pi)**(n*p/2) * multigamma((nu + p - 1)/2, p)) + return K * (Determinant(eye(n) + Inverse(Sigma)*(x - M)*Inverse(Omega)*Transpose(x - M))) \ + **(-(nu + n + p -1)/2) + + + +def MatrixStudentT(symbol, nu, location_matrix, scale_matrix_1, scale_matrix_2): + """ + Creates a random variable with Matrix Gamma Distribution. + + The density of the said distribution can be found at [1]. + + Parameters + ========== + + nu: Positive Real number + degrees of freedom + location_matrix: Positive definite real square matrix + Location Matrix of shape ``n x p`` + scale_matrix_1: Positive definite real square matrix + Scale Matrix of shape ``p x p`` + scale_matrix_2: Positive definite real square matrix + Scale Matrix of shape ``n x n`` + + Returns + ======= + + RandomSymbol + + Examples + ======== + + >>> from sympy import MatrixSymbol,symbols + >>> from sympy.stats import density, MatrixStudentT + >>> v = symbols('v',positive=True) + >>> M = MatrixStudentT('M', v, [[1, 2]], [[1, 0], [0, 1]], [1]) + >>> X = MatrixSymbol('X', 1, 2) + >>> density(M)(X) + gamma(v/2 + 1)*Determinant((Matrix([[-1, -2]]) + X)*(Matrix([ + [-1], + [-2]]) + X.T) + Matrix([[1]]))**(-v/2 - 1)/(pi**1.0*gamma(v/2)*Determinant(Matrix([[1]]))**1.0*Determinant(Matrix([ + [1, 0], + [0, 1]]))**0.5) + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Matrix_t-distribution + + """ + if isinstance(location_matrix, list): + location_matrix = ImmutableMatrix(location_matrix) + if isinstance(scale_matrix_1, list): + scale_matrix_1 = ImmutableMatrix(scale_matrix_1) + if isinstance(scale_matrix_2, list): + scale_matrix_2 = ImmutableMatrix(scale_matrix_2) + args = (nu, location_matrix, scale_matrix_1, scale_matrix_2) + return rv(symbol, MatrixStudentTDistribution, args) diff --git a/lib/python3.12/site-packages/sympy/stats/random_matrix.py b/lib/python3.12/site-packages/sympy/stats/random_matrix.py new file mode 100644 index 0000000000000000000000000000000000000000..fdd25cb9ad23fed9d3a85982b24bef33d04928f0 --- /dev/null +++ b/lib/python3.12/site-packages/sympy/stats/random_matrix.py @@ -0,0 +1,30 @@ +from sympy.core.basic import Basic +from sympy.stats.rv import PSpace, _symbol_converter, RandomMatrixSymbol + +class RandomMatrixPSpace(PSpace): + """ + Represents probability space for + random matrices. It contains the mechanics + for handling the API calls for random matrices. + """ + def __new__(cls, sym, model=None): + sym = _symbol_converter(sym) + if model: + return Basic.__new__(cls, sym, model) + else: + return Basic.__new__(cls, sym) + + @property + def model(self): + try: + return self.args[1] + except IndexError: + return None + + def compute_density(self, expr, *args): + rms = expr.atoms(RandomMatrixSymbol) + if len(rms) > 2 or (not isinstance(expr, RandomMatrixSymbol)): + raise NotImplementedError("Currently, no algorithm has been " + "implemented to handle general expressions containing " + "multiple random matrices.") + return self.model.density(expr) diff --git a/lib/python3.12/site-packages/sympy/stats/random_matrix_models.py b/lib/python3.12/site-packages/sympy/stats/random_matrix_models.py new file mode 100644 index 0000000000000000000000000000000000000000..6327a248ea5919c0bbb0ffc2c984105e04fe20e9 --- /dev/null +++ b/lib/python3.12/site-packages/sympy/stats/random_matrix_models.py @@ -0,0 +1,457 @@ +from sympy.concrete.products import Product +from sympy.concrete.summations import Sum +from sympy.core.basic import Basic +from sympy.core.function import Lambda +from sympy.core.numbers import (I, pi) +from sympy.core.singleton import S +from sympy.core.symbol import Dummy +from sympy.functions.elementary.complexes import Abs +from sympy.functions.elementary.exponential import exp +from sympy.functions.special.gamma_functions import gamma +from sympy.integrals.integrals import Integral +from sympy.matrices.expressions.matexpr import MatrixSymbol +from sympy.matrices.expressions.trace import Trace +from sympy.tensor.indexed import IndexedBase +from sympy.core.sympify import _sympify +from sympy.stats.rv import _symbol_converter, Density, RandomMatrixSymbol, is_random +from sympy.stats.joint_rv_types import JointDistributionHandmade +from sympy.stats.random_matrix import RandomMatrixPSpace +from sympy.tensor.array import ArrayComprehension + +__all__ = [ + 'CircularEnsemble', + 'CircularUnitaryEnsemble', + 'CircularOrthogonalEnsemble', + 'CircularSymplecticEnsemble', + 'GaussianEnsemble', + 'GaussianUnitaryEnsemble', + 'GaussianOrthogonalEnsemble', + 'GaussianSymplecticEnsemble', + 'joint_eigen_distribution', + 'JointEigenDistribution', + 'level_spacing_distribution' +] + +@is_random.register(RandomMatrixSymbol) +def _(x): + return True + + +class RandomMatrixEnsembleModel(Basic): + """ + Base class for random matrix ensembles. + It acts as an umbrella and contains + the methods common to all the ensembles + defined in sympy.stats.random_matrix_models. + """ + def __new__(cls, sym, dim=None): + sym, dim = _symbol_converter(sym), _sympify(dim) + if dim.is_integer == False: + raise ValueError("Dimension of the random matrices must be " + "integers, received %s instead."%(dim)) + return Basic.__new__(cls, sym, dim) + + symbol = property(lambda self: self.args[0]) + dimension = property(lambda self: self.args[1]) + + def density(self, expr): + return Density(expr) + + def __call__(self, expr): + return self.density(expr) + +class GaussianEnsembleModel(RandomMatrixEnsembleModel): + """ + Abstract class for Gaussian ensembles. + Contains the properties common to all the + gaussian ensembles. + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Random_matrix#Gaussian_ensembles + .. [2] https://arxiv.org/pdf/1712.07903.pdf + """ + def _compute_normalization_constant(self, beta, n): + """ + Helper function for computing normalization + constant for joint probability density of eigen + values of Gaussian ensembles. + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Selberg_integral#Mehta's_integral + """ + n = S(n) + prod_term = lambda j: gamma(1 + beta*S(j)/2)/gamma(S.One + beta/S(2)) + j = Dummy('j', integer=True, positive=True) + term1 = Product(prod_term(j), (j, 1, n)).doit() + term2 = (2/(beta*n))**(beta*n*(n - 1)/4 + n/2) + term3 = (2*pi)**(n/2) + return term1 * term2 * term3 + + def _compute_joint_eigen_distribution(self, beta): + """ + Helper function for computing the joint + probability distribution of eigen values + of the random matrix. + """ + n = self.dimension + Zbn = self._compute_normalization_constant(beta, n) + l = IndexedBase('l') + i = Dummy('i', integer=True, positive=True) + j = Dummy('j', integer=True, positive=True) + k = Dummy('k', integer=True, positive=True) + term1 = exp((-S(n)/2) * Sum(l[k]**2, (k, 1, n)).doit()) + sub_term = Lambda(i, Product(Abs(l[j] - l[i])**beta, (j, i + 1, n))) + term2 = Product(sub_term(i).doit(), (i, 1, n - 1)).doit() + syms = ArrayComprehension(l[k], (k, 1, n)).doit() + return Lambda(tuple(syms), (term1 * term2)/Zbn) + +class GaussianUnitaryEnsembleModel(GaussianEnsembleModel): + @property + def normalization_constant(self): + n = self.dimension + return 2**(S(n)/2) * pi**(S(n**2)/2) + + def density(self, expr): + n, ZGUE = self.dimension, self.normalization_constant + h_pspace = RandomMatrixPSpace('P', model=self) + H = RandomMatrixSymbol('H', n, n, pspace=h_pspace) + return Lambda(H, exp(-S(n)/2 * Trace(H**2))/ZGUE)(expr) + + def joint_eigen_distribution(self): + return self._compute_joint_eigen_distribution(S(2)) + + def level_spacing_distribution(self): + s = Dummy('s') + f = (32/pi**2)*(s**2)*exp((-4/pi)*s**2) + return Lambda(s, f) + +class GaussianOrthogonalEnsembleModel(GaussianEnsembleModel): + @property + def normalization_constant(self): + n = self.dimension + _H = MatrixSymbol('_H', n, n) + return Integral(exp(-S(n)/4 * Trace(_H**2))) + + def density(self, expr): + n, ZGOE = self.dimension, self.normalization_constant + h_pspace = RandomMatrixPSpace('P', model=self) + H = RandomMatrixSymbol('H', n, n, pspace=h_pspace) + return Lambda(H, exp(-S(n)/4 * Trace(H**2))/ZGOE)(expr) + + def joint_eigen_distribution(self): + return self._compute_joint_eigen_distribution(S.One) + + def level_spacing_distribution(self): + s = Dummy('s') + f = (pi/2)*s*exp((-pi/4)*s**2) + return Lambda(s, f) + +class GaussianSymplecticEnsembleModel(GaussianEnsembleModel): + @property + def normalization_constant(self): + n = self.dimension + _H = MatrixSymbol('_H', n, n) + return Integral(exp(-S(n) * Trace(_H**2))) + + def density(self, expr): + n, ZGSE = self.dimension, self.normalization_constant + h_pspace = RandomMatrixPSpace('P', model=self) + H = RandomMatrixSymbol('H', n, n, pspace=h_pspace) + return Lambda(H, exp(-S(n) * Trace(H**2))/ZGSE)(expr) + + def joint_eigen_distribution(self): + return self._compute_joint_eigen_distribution(S(4)) + + def level_spacing_distribution(self): + s = Dummy('s') + f = ((S(2)**18)/((S(3)**6)*(pi**3)))*(s**4)*exp((-64/(9*pi))*s**2) + return Lambda(s, f) + +def GaussianEnsemble(sym, dim): + sym, dim = _symbol_converter(sym), _sympify(dim) + model = GaussianEnsembleModel(sym, dim) + rmp = RandomMatrixPSpace(sym, model=model) + return RandomMatrixSymbol(sym, dim, dim, pspace=rmp) + +def GaussianUnitaryEnsemble(sym, dim): + """ + Represents Gaussian Unitary Ensembles. + + Examples + ======== + + >>> from sympy.stats import GaussianUnitaryEnsemble as GUE, density + >>> from sympy import MatrixSymbol + >>> G = GUE('U', 2) + >>> X = MatrixSymbol('X', 2, 2) + >>> density(G)(X) + exp(-Trace(X**2))/(2*pi**2) + """ + sym, dim = _symbol_converter(sym), _sympify(dim) + model = GaussianUnitaryEnsembleModel(sym, dim) + rmp = RandomMatrixPSpace(sym, model=model) + return RandomMatrixSymbol(sym, dim, dim, pspace=rmp) + +def GaussianOrthogonalEnsemble(sym, dim): + """ + Represents Gaussian Orthogonal Ensembles. + + Examples + ======== + + >>> from sympy.stats import GaussianOrthogonalEnsemble as GOE, density + >>> from sympy import MatrixSymbol + >>> G = GOE('U', 2) + >>> X = MatrixSymbol('X', 2, 2) + >>> density(G)(X) + exp(-Trace(X**2)/2)/Integral(exp(-Trace(_H**2)/2), _H) + """ + sym, dim = _symbol_converter(sym), _sympify(dim) + model = GaussianOrthogonalEnsembleModel(sym, dim) + rmp = RandomMatrixPSpace(sym, model=model) + return RandomMatrixSymbol(sym, dim, dim, pspace=rmp) + +def GaussianSymplecticEnsemble(sym, dim): + """ + Represents Gaussian Symplectic Ensembles. + + Examples + ======== + + >>> from sympy.stats import GaussianSymplecticEnsemble as GSE, density + >>> from sympy import MatrixSymbol + >>> G = GSE('U', 2) + >>> X = MatrixSymbol('X', 2, 2) + >>> density(G)(X) + exp(-2*Trace(X**2))/Integral(exp(-2*Trace(_H**2)), _H) + """ + sym, dim = _symbol_converter(sym), _sympify(dim) + model = GaussianSymplecticEnsembleModel(sym, dim) + rmp = RandomMatrixPSpace(sym, model=model) + return RandomMatrixSymbol(sym, dim, dim, pspace=rmp) + +class CircularEnsembleModel(RandomMatrixEnsembleModel): + """ + Abstract class for Circular ensembles. + Contains the properties and methods + common to all the circular ensembles. + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Circular_ensemble + """ + def density(self, expr): + # TODO : Add support for Lie groups(as extensions of sympy.diffgeom) + # and define measures on them + raise NotImplementedError("Support for Haar measure hasn't been " + "implemented yet, therefore the density of " + "%s cannot be computed."%(self)) + + def _compute_joint_eigen_distribution(self, beta): + """ + Helper function to compute the joint distribution of phases + of the complex eigen values of matrices belonging to any + circular ensembles. + """ + n = self.dimension + Zbn = ((2*pi)**n)*(gamma(beta*n/2 + 1)/S(gamma(beta/2 + 1))**n) + t = IndexedBase('t') + i, j, k = (Dummy('i', integer=True), Dummy('j', integer=True), + Dummy('k', integer=True)) + syms = ArrayComprehension(t[i], (i, 1, n)).doit() + f = Product(Product(Abs(exp(I*t[k]) - exp(I*t[j]))**beta, (j, k + 1, n)).doit(), + (k, 1, n - 1)).doit() + return Lambda(tuple(syms), f/Zbn) + +class CircularUnitaryEnsembleModel(CircularEnsembleModel): + def joint_eigen_distribution(self): + return self._compute_joint_eigen_distribution(S(2)) + +class CircularOrthogonalEnsembleModel(CircularEnsembleModel): + def joint_eigen_distribution(self): + return self._compute_joint_eigen_distribution(S.One) + +class CircularSymplecticEnsembleModel(CircularEnsembleModel): + def joint_eigen_distribution(self): + return self._compute_joint_eigen_distribution(S(4)) + +def CircularEnsemble(sym, dim): + sym, dim = _symbol_converter(sym), _sympify(dim) + model = CircularEnsembleModel(sym, dim) + rmp = RandomMatrixPSpace(sym, model=model) + return RandomMatrixSymbol(sym, dim, dim, pspace=rmp) + +def CircularUnitaryEnsemble(sym, dim): + """ + Represents Circular Unitary Ensembles. + + Examples + ======== + + >>> from sympy.stats import CircularUnitaryEnsemble as CUE + >>> from sympy.stats import joint_eigen_distribution + >>> C = CUE('U', 1) + >>> joint_eigen_distribution(C) + Lambda(t[1], Product(Abs(exp(I*t[_j]) - exp(I*t[_k]))**2, (_j, _k + 1, 1), (_k, 1, 0))/(2*pi)) + + Note + ==== + + As can be seen above in the example, density of CiruclarUnitaryEnsemble + is not evaluated because the exact definition is based on haar measure of + unitary group which is not unique. + """ + sym, dim = _symbol_converter(sym), _sympify(dim) + model = CircularUnitaryEnsembleModel(sym, dim) + rmp = RandomMatrixPSpace(sym, model=model) + return RandomMatrixSymbol(sym, dim, dim, pspace=rmp) + +def CircularOrthogonalEnsemble(sym, dim): + """ + Represents Circular Orthogonal Ensembles. + + Examples + ======== + + >>> from sympy.stats import CircularOrthogonalEnsemble as COE + >>> from sympy.stats import joint_eigen_distribution + >>> C = COE('O', 1) + >>> joint_eigen_distribution(C) + Lambda(t[1], Product(Abs(exp(I*t[_j]) - exp(I*t[_k])), (_j, _k + 1, 1), (_k, 1, 0))/(2*pi)) + + Note + ==== + + As can be seen above in the example, density of CiruclarOrthogonalEnsemble + is not evaluated because the exact definition is based on haar measure of + unitary group which is not unique. + """ + sym, dim = _symbol_converter(sym), _sympify(dim) + model = CircularOrthogonalEnsembleModel(sym, dim) + rmp = RandomMatrixPSpace(sym, model=model) + return RandomMatrixSymbol(sym, dim, dim, pspace=rmp) + +def CircularSymplecticEnsemble(sym, dim): + """ + Represents Circular Symplectic Ensembles. + + Examples + ======== + + >>> from sympy.stats import CircularSymplecticEnsemble as CSE + >>> from sympy.stats import joint_eigen_distribution + >>> C = CSE('S', 1) + >>> joint_eigen_distribution(C) + Lambda(t[1], Product(Abs(exp(I*t[_j]) - exp(I*t[_k]))**4, (_j, _k + 1, 1), (_k, 1, 0))/(2*pi)) + + Note + ==== + + As can be seen above in the example, density of CiruclarSymplecticEnsemble + is not evaluated because the exact definition is based on haar measure of + unitary group which is not unique. + """ + sym, dim = _symbol_converter(sym), _sympify(dim) + model = CircularSymplecticEnsembleModel(sym, dim) + rmp = RandomMatrixPSpace(sym, model=model) + return RandomMatrixSymbol(sym, dim, dim, pspace=rmp) + +def joint_eigen_distribution(mat): + """ + For obtaining joint probability distribution + of eigen values of random matrix. + + Parameters + ========== + + mat: RandomMatrixSymbol + The matrix symbol whose eigen values are to be considered. + + Returns + ======= + + Lambda + + Examples + ======== + + >>> from sympy.stats import GaussianUnitaryEnsemble as GUE + >>> from sympy.stats import joint_eigen_distribution + >>> U = GUE('U', 2) + >>> joint_eigen_distribution(U) + Lambda((l[1], l[2]), exp(-l[1]**2 - l[2]**2)*Product(Abs(l[_i] - l[_j])**2, (_j, _i + 1, 2), (_i, 1, 1))/pi) + """ + if not isinstance(mat, RandomMatrixSymbol): + raise ValueError("%s is not of type, RandomMatrixSymbol."%(mat)) + return mat.pspace.model.joint_eigen_distribution() + +def JointEigenDistribution(mat): + """ + Creates joint distribution of eigen values of matrices with random + expressions. + + Parameters + ========== + + mat: Matrix + The matrix under consideration. + + Returns + ======= + + JointDistributionHandmade + + Examples + ======== + + >>> from sympy.stats import Normal, JointEigenDistribution + >>> from sympy import Matrix + >>> A = [[Normal('A00', 0, 1), Normal('A01', 0, 1)], + ... [Normal('A10', 0, 1), Normal('A11', 0, 1)]] + >>> JointEigenDistribution(Matrix(A)) + JointDistributionHandmade(-sqrt(A00**2 - 2*A00*A11 + 4*A01*A10 + A11**2)/2 + + A00/2 + A11/2, sqrt(A00**2 - 2*A00*A11 + 4*A01*A10 + A11**2)/2 + A00/2 + A11/2) + + """ + eigenvals = mat.eigenvals(multiple=True) + if not all(is_random(eigenval) for eigenval in set(eigenvals)): + raise ValueError("Eigen values do not have any random expression, " + "joint distribution cannot be generated.") + return JointDistributionHandmade(*eigenvals) + +def level_spacing_distribution(mat): + """ + For obtaining distribution of level spacings. + + Parameters + ========== + + mat: RandomMatrixSymbol + The random matrix symbol whose eigen values are + to be considered for finding the level spacings. + + Returns + ======= + + Lambda + + Examples + ======== + + >>> from sympy.stats import GaussianUnitaryEnsemble as GUE + >>> from sympy.stats import level_spacing_distribution + >>> U = GUE('U', 2) + >>> level_spacing_distribution(U) + Lambda(_s, 32*_s**2*exp(-4*_s**2/pi)/pi**2) + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Random_matrix#Distribution_of_level_spacings + """ + return mat.pspace.model.level_spacing_distribution() diff --git a/lib/python3.12/site-packages/sympy/stats/rv.py b/lib/python3.12/site-packages/sympy/stats/rv.py new file mode 100644 index 0000000000000000000000000000000000000000..2344ae9397e46b27f6ae178ea9a4030c747adbf3 --- /dev/null +++ b/lib/python3.12/site-packages/sympy/stats/rv.py @@ -0,0 +1,1798 @@ +""" +Main Random Variables Module + +Defines abstract random variable type. +Contains interfaces for probability space object (PSpace) as well as standard +operators, P, E, sample, density, where, quantile + +See Also +======== + +sympy.stats.crv +sympy.stats.frv +sympy.stats.rv_interface +""" + +from __future__ import annotations +from functools import singledispatch +from math import prod + +from sympy.core.add import Add +from sympy.core.basic import Basic +from sympy.core.containers import Tuple +from sympy.core.expr import Expr +from sympy.core.function import (Function, Lambda) +from sympy.core.logic import fuzzy_and +from sympy.core.mul import Mul +from sympy.core.relational import (Eq, Ne) +from sympy.core.singleton import S +from sympy.core.symbol import (Dummy, Symbol) +from sympy.core.sympify import sympify +from sympy.functions.special.delta_functions import DiracDelta +from sympy.functions.special.tensor_functions import KroneckerDelta +from sympy.logic.boolalg import (And, Or) +from sympy.matrices.expressions.matexpr import MatrixSymbol +from sympy.tensor.indexed import Indexed +from sympy.utilities.lambdify import lambdify +from sympy.core.relational import Relational +from sympy.core.sympify import _sympify +from sympy.sets.sets import FiniteSet, ProductSet, Intersection +from sympy.solvers.solveset import solveset +from sympy.external import import_module +from sympy.utilities.decorator import doctest_depends_on +from sympy.utilities.exceptions import sympy_deprecation_warning +from sympy.utilities.iterables import iterable + + +__doctest_requires__ = {('sample',): ['scipy']} + + +x = Symbol('x') + + +@singledispatch +def is_random(x): + return False + + +@is_random.register(Basic) +def _(x): + atoms = x.free_symbols + return any(is_random(i) for i in atoms) + + +class RandomDomain(Basic): + """ + Represents a set of variables and the values which they can take. + + See Also + ======== + + sympy.stats.crv.ContinuousDomain + sympy.stats.frv.FiniteDomain + """ + + is_ProductDomain = False + is_Finite = False + is_Continuous = False + is_Discrete = False + + def __new__(cls, symbols, *args): + symbols = FiniteSet(*symbols) + return Basic.__new__(cls, symbols, *args) + + @property + def symbols(self): + return self.args[0] + + @property + def set(self): + return self.args[1] + + def __contains__(self, other): + raise NotImplementedError() + + def compute_expectation(self, expr): + raise NotImplementedError() + + +class SingleDomain(RandomDomain): + """ + A single variable and its domain. + + See Also + ======== + + sympy.stats.crv.SingleContinuousDomain + sympy.stats.frv.SingleFiniteDomain + """ + def __new__(cls, symbol, set): + assert symbol.is_Symbol + return Basic.__new__(cls, symbol, set) + + @property + def symbol(self): + return self.args[0] + + @property + def symbols(self): + return FiniteSet(self.symbol) + + def __contains__(self, other): + if len(other) != 1: + return False + sym, val = tuple(other)[0] + return self.symbol == sym and val in self.set + + +class MatrixDomain(RandomDomain): + """ + A Random Matrix variable and its domain. + + """ + def __new__(cls, symbol, set): + symbol, set = _symbol_converter(symbol), _sympify(set) + return Basic.__new__(cls, symbol, set) + + @property + def symbol(self): + return self.args[0] + + @property + def symbols(self): + return FiniteSet(self.symbol) + + +class ConditionalDomain(RandomDomain): + """ + A RandomDomain with an attached condition. + + See Also + ======== + + sympy.stats.crv.ConditionalContinuousDomain + sympy.stats.frv.ConditionalFiniteDomain + """ + def __new__(cls, fulldomain, condition): + condition = condition.xreplace({rs: rs.symbol + for rs in random_symbols(condition)}) + return Basic.__new__(cls, fulldomain, condition) + + @property + def symbols(self): + return self.fulldomain.symbols + + @property + def fulldomain(self): + return self.args[0] + + @property + def condition(self): + return self.args[1] + + @property + def set(self): + raise NotImplementedError("Set of Conditional Domain not Implemented") + + def as_boolean(self): + return And(self.fulldomain.as_boolean(), self.condition) + + +class PSpace(Basic): + """ + A Probability Space. + + Explanation + =========== + + Probability Spaces encode processes that equal different values + probabilistically. These underly Random Symbols which occur in SymPy + expressions and contain the mechanics to evaluate statistical statements. + + See Also + ======== + + sympy.stats.crv.ContinuousPSpace + sympy.stats.frv.FinitePSpace + """ + + is_Finite = None # type: bool + is_Continuous = None # type: bool + is_Discrete = None # type: bool + is_real = None # type: bool + + @property + def domain(self): + return self.args[0] + + @property + def density(self): + return self.args[1] + + @property + def values(self): + return frozenset(RandomSymbol(sym, self) for sym in self.symbols) + + @property + def symbols(self): + return self.domain.symbols + + def where(self, condition): + raise NotImplementedError() + + def compute_density(self, expr): + raise NotImplementedError() + + def sample(self, size=(), library='scipy', seed=None): + raise NotImplementedError() + + def probability(self, condition): + raise NotImplementedError() + + def compute_expectation(self, expr): + raise NotImplementedError() + + +class SinglePSpace(PSpace): + """ + Represents the probabilities of a set of random events that can be + attributed to a single variable/symbol. + """ + def __new__(cls, s, distribution): + s = _symbol_converter(s) + return Basic.__new__(cls, s, distribution) + + @property + def value(self): + return RandomSymbol(self.symbol, self) + + @property + def symbol(self): + return self.args[0] + + @property + def distribution(self): + return self.args[1] + + @property + def pdf(self): + return self.distribution.pdf(self.symbol) + + +class RandomSymbol(Expr): + """ + Random Symbols represent ProbabilitySpaces in SymPy Expressions. + In principle they can take on any value that their symbol can take on + within the associated PSpace with probability determined by the PSpace + Density. + + Explanation + =========== + + Random Symbols contain pspace and symbol properties. + The pspace property points to the represented Probability Space + The symbol is a standard SymPy Symbol that is used in that probability space + for example in defining a density. + + You can form normal SymPy expressions using RandomSymbols and operate on + those expressions with the Functions + + E - Expectation of a random expression + P - Probability of a condition + density - Probability Density of an expression + given - A new random expression (with new random symbols) given a condition + + An object of the RandomSymbol type should almost never be created by the + user. They tend to be created instead by the PSpace class's value method. + Traditionally a user does not even do this but instead calls one of the + convenience functions Normal, Exponential, Coin, Die, FiniteRV, etc.... + """ + + def __new__(cls, symbol, pspace=None): + from sympy.stats.joint_rv import JointRandomSymbol + if pspace is None: + # Allow single arg, representing pspace == PSpace() + pspace = PSpace() + symbol = _symbol_converter(symbol) + if not isinstance(pspace, PSpace): + raise TypeError("pspace variable should be of type PSpace") + if cls == JointRandomSymbol and isinstance(pspace, SinglePSpace): + cls = RandomSymbol + return Basic.__new__(cls, symbol, pspace) + + is_finite = True + is_symbol = True + is_Atom = True + + _diff_wrt = True + + pspace = property(lambda self: self.args[1]) + symbol = property(lambda self: self.args[0]) + name = property(lambda self: self.symbol.name) + + def _eval_is_positive(self): + return self.symbol.is_positive + + def _eval_is_integer(self): + return self.symbol.is_integer + + def _eval_is_real(self): + return self.symbol.is_real or self.pspace.is_real + + @property + def is_commutative(self): + return self.symbol.is_commutative + + @property + def free_symbols(self): + return {self} + +class RandomIndexedSymbol(RandomSymbol): + + def __new__(cls, idx_obj, pspace=None): + if pspace is None: + # Allow single arg, representing pspace == PSpace() + pspace = PSpace() + if not isinstance(idx_obj, (Indexed, Function)): + raise TypeError("An Function or Indexed object is expected not %s"%(idx_obj)) + return Basic.__new__(cls, idx_obj, pspace) + + symbol = property(lambda self: self.args[0]) + name = property(lambda self: str(self.args[0])) + + @property + def key(self): + if isinstance(self.symbol, Indexed): + return self.symbol.args[1] + elif isinstance(self.symbol, Function): + return self.symbol.args[0] + + @property + def free_symbols(self): + if self.key.free_symbols: + free_syms = self.key.free_symbols + free_syms.add(self) + return free_syms + return {self} + + @property + def pspace(self): + return self.args[1] + +class RandomMatrixSymbol(RandomSymbol, MatrixSymbol): # type: ignore + def __new__(cls, symbol, n, m, pspace=None): + n, m = _sympify(n), _sympify(m) + symbol = _symbol_converter(symbol) + if pspace is None: + # Allow single arg, representing pspace == PSpace() + pspace = PSpace() + return Basic.__new__(cls, symbol, n, m, pspace) + + symbol = property(lambda self: self.args[0]) + pspace = property(lambda self: self.args[3]) + + +class ProductPSpace(PSpace): + """ + Abstract class for representing probability spaces with multiple random + variables. + + See Also + ======== + + sympy.stats.rv.IndependentProductPSpace + sympy.stats.joint_rv.JointPSpace + """ + pass + +class IndependentProductPSpace(ProductPSpace): + """ + A probability space resulting from the merger of two independent probability + spaces. + + Often created using the function, pspace. + """ + + def __new__(cls, *spaces): + rs_space_dict = {} + for space in spaces: + for value in space.values: + rs_space_dict[value] = space + + symbols = FiniteSet(*[val.symbol for val in rs_space_dict.keys()]) + + # Overlapping symbols + from sympy.stats.joint_rv import MarginalDistribution + from sympy.stats.compound_rv import CompoundDistribution + if len(symbols) < sum(len(space.symbols) for space in spaces if not + isinstance(space.distribution, ( + CompoundDistribution, MarginalDistribution))): + raise ValueError("Overlapping Random Variables") + + if all(space.is_Finite for space in spaces): + from sympy.stats.frv import ProductFinitePSpace + cls = ProductFinitePSpace + + obj = Basic.__new__(cls, *FiniteSet(*spaces)) + + return obj + + @property + def pdf(self): + p = Mul(*[space.pdf for space in self.spaces]) + return p.subs({rv: rv.symbol for rv in self.values}) + + @property + def rs_space_dict(self): + d = {} + for space in self.spaces: + for value in space.values: + d[value] = space + return d + + @property + def symbols(self): + return FiniteSet(*[val.symbol for val in self.rs_space_dict.keys()]) + + @property + def spaces(self): + return FiniteSet(*self.args) + + @property + def values(self): + return sumsets(space.values for space in self.spaces) + + def compute_expectation(self, expr, rvs=None, evaluate=False, **kwargs): + rvs = rvs or self.values + rvs = frozenset(rvs) + for space in self.spaces: + expr = space.compute_expectation(expr, rvs & space.values, evaluate=False, **kwargs) + if evaluate and hasattr(expr, 'doit'): + return expr.doit(**kwargs) + return expr + + @property + def domain(self): + return ProductDomain(*[space.domain for space in self.spaces]) + + @property + def density(self): + raise NotImplementedError("Density not available for ProductSpaces") + + def sample(self, size=(), library='scipy', seed=None): + return {k: v for space in self.spaces + for k, v in space.sample(size=size, library=library, seed=seed).items()} + + + def probability(self, condition, **kwargs): + cond_inv = False + if isinstance(condition, Ne): + condition = Eq(condition.args[0], condition.args[1]) + cond_inv = True + elif isinstance(condition, And): # they are independent + return Mul(*[self.probability(arg) for arg in condition.args]) + elif isinstance(condition, Or): # they are independent + return Add(*[self.probability(arg) for arg in condition.args]) + expr = condition.lhs - condition.rhs + rvs = random_symbols(expr) + dens = self.compute_density(expr) + if any(pspace(rv).is_Continuous for rv in rvs): + from sympy.stats.crv import SingleContinuousPSpace + from sympy.stats.crv_types import ContinuousDistributionHandmade + if expr in self.values: + # Marginalize all other random symbols out of the density + randomsymbols = tuple(set(self.values) - frozenset([expr])) + symbols = tuple(rs.symbol for rs in randomsymbols) + pdf = self.domain.integrate(self.pdf, symbols, **kwargs) + return Lambda(expr.symbol, pdf) + dens = ContinuousDistributionHandmade(dens) + z = Dummy('z', real=True) + space = SingleContinuousPSpace(z, dens) + result = space.probability(condition.__class__(space.value, 0)) + else: + from sympy.stats.drv import SingleDiscretePSpace + from sympy.stats.drv_types import DiscreteDistributionHandmade + dens = DiscreteDistributionHandmade(dens) + z = Dummy('z', integer=True) + space = SingleDiscretePSpace(z, dens) + result = space.probability(condition.__class__(space.value, 0)) + return result if not cond_inv else S.One - result + + def compute_density(self, expr, **kwargs): + rvs = random_symbols(expr) + if any(pspace(rv).is_Continuous for rv in rvs): + z = Dummy('z', real=True) + expr = self.compute_expectation(DiracDelta(expr - z), + **kwargs) + else: + z = Dummy('z', integer=True) + expr = self.compute_expectation(KroneckerDelta(expr, z), + **kwargs) + return Lambda(z, expr) + + def compute_cdf(self, expr, **kwargs): + raise ValueError("CDF not well defined on multivariate expressions") + + def conditional_space(self, condition, normalize=True, **kwargs): + rvs = random_symbols(condition) + condition = condition.xreplace({rv: rv.symbol for rv in self.values}) + pspaces = [pspace(rv) for rv in rvs] + if any(ps.is_Continuous for ps in pspaces): + from sympy.stats.crv import (ConditionalContinuousDomain, + ContinuousPSpace) + space = ContinuousPSpace + domain = ConditionalContinuousDomain(self.domain, condition) + elif any(ps.is_Discrete for ps in pspaces): + from sympy.stats.drv import (ConditionalDiscreteDomain, + DiscretePSpace) + space = DiscretePSpace + domain = ConditionalDiscreteDomain(self.domain, condition) + elif all(ps.is_Finite for ps in pspaces): + from sympy.stats.frv import FinitePSpace + return FinitePSpace.conditional_space(self, condition) + if normalize: + replacement = {rv: Dummy(str(rv)) for rv in self.symbols} + norm = domain.compute_expectation(self.pdf, **kwargs) + pdf = self.pdf / norm.xreplace(replacement) + # XXX: Converting symbols from set to tuple. The order matters to + # Lambda though so we shouldn't be starting with a set here... + density = Lambda(tuple(domain.symbols), pdf) + + return space(domain, density) + +class ProductDomain(RandomDomain): + """ + A domain resulting from the merger of two independent domains. + + See Also + ======== + sympy.stats.crv.ProductContinuousDomain + sympy.stats.frv.ProductFiniteDomain + """ + is_ProductDomain = True + + def __new__(cls, *domains): + # Flatten any product of products + domains2 = [] + for domain in domains: + if not domain.is_ProductDomain: + domains2.append(domain) + else: + domains2.extend(domain.domains) + domains2 = FiniteSet(*domains2) + + if all(domain.is_Finite for domain in domains2): + from sympy.stats.frv import ProductFiniteDomain + cls = ProductFiniteDomain + if all(domain.is_Continuous for domain in domains2): + from sympy.stats.crv import ProductContinuousDomain + cls = ProductContinuousDomain + if all(domain.is_Discrete for domain in domains2): + from sympy.stats.drv import ProductDiscreteDomain + cls = ProductDiscreteDomain + + return Basic.__new__(cls, *domains2) + + @property + def sym_domain_dict(self): + return {symbol: domain for domain in self.domains + for symbol in domain.symbols} + + @property + def symbols(self): + return FiniteSet(*[sym for domain in self.domains + for sym in domain.symbols]) + + @property + def domains(self): + return self.args + + @property + def set(self): + return ProductSet(*(domain.set for domain in self.domains)) + + def __contains__(self, other): + # Split event into each subdomain + for domain in self.domains: + # Collect the parts of this event which associate to this domain + elem = frozenset([item for item in other + if sympify(domain.symbols.contains(item[0])) + is S.true]) + # Test this sub-event + if elem not in domain: + return False + # All subevents passed + return True + + def as_boolean(self): + return And(*[domain.as_boolean() for domain in self.domains]) + + +def random_symbols(expr): + """ + Returns all RandomSymbols within a SymPy Expression. + """ + atoms = getattr(expr, 'atoms', None) + if atoms is not None: + comp = lambda rv: rv.symbol.name + l = list(atoms(RandomSymbol)) + return sorted(l, key=comp) + else: + return [] + + +def pspace(expr): + """ + Returns the underlying Probability Space of a random expression. + + For internal use. + + Examples + ======== + + >>> from sympy.stats import pspace, Normal + >>> X = Normal('X', 0, 1) + >>> pspace(2*X + 1) == X.pspace + True + """ + expr = sympify(expr) + if isinstance(expr, RandomSymbol) and expr.pspace is not None: + return expr.pspace + if expr.has(RandomMatrixSymbol): + rm = list(expr.atoms(RandomMatrixSymbol))[0] + return rm.pspace + + rvs = random_symbols(expr) + if not rvs: + raise ValueError("Expression containing Random Variable expected, not %s" % (expr)) + # If only one space present + if all(rv.pspace == rvs[0].pspace for rv in rvs): + return rvs[0].pspace + from sympy.stats.compound_rv import CompoundPSpace + from sympy.stats.stochastic_process import StochasticPSpace + for rv in rvs: + if isinstance(rv.pspace, (CompoundPSpace, StochasticPSpace)): + return rv.pspace + # Otherwise make a product space + return IndependentProductPSpace(*[rv.pspace for rv in rvs]) + + +def sumsets(sets): + """ + Union of sets + """ + return frozenset().union(*sets) + + +def rs_swap(a, b): + """ + Build a dictionary to swap RandomSymbols based on their underlying symbol. + + i.e. + if ``X = ('x', pspace1)`` + and ``Y = ('x', pspace2)`` + then ``X`` and ``Y`` match and the key, value pair + ``{X:Y}`` will appear in the result + + Inputs: collections a and b of random variables which share common symbols + Output: dict mapping RVs in a to RVs in b + """ + d = {} + for rsa in a: + d[rsa] = [rsb for rsb in b if rsa.symbol == rsb.symbol][0] + return d + + +def given(expr, condition=None, **kwargs): + r""" Conditional Random Expression. + + Explanation + =========== + + From a random expression and a condition on that expression creates a new + probability space from the condition and returns the same expression on that + conditional probability space. + + Examples + ======== + + >>> from sympy.stats import given, density, Die + >>> X = Die('X', 6) + >>> Y = given(X, X > 3) + >>> density(Y).dict + {4: 1/3, 5: 1/3, 6: 1/3} + + Following convention, if the condition is a random symbol then that symbol + is considered fixed. + + >>> from sympy.stats import Normal + >>> from sympy import pprint + >>> from sympy.abc import z + + >>> X = Normal('X', 0, 1) + >>> Y = Normal('Y', 0, 1) + >>> pprint(density(X + Y, Y)(z), use_unicode=False) + 2 + -(-Y + z) + ----------- + ___ 2 + \/ 2 *e + ------------------ + ____ + 2*\/ pi + """ + + if not is_random(condition) or pspace_independent(expr, condition): + return expr + + if isinstance(condition, RandomSymbol): + condition = Eq(condition, condition.symbol) + + condsymbols = random_symbols(condition) + if (isinstance(condition, Eq) and len(condsymbols) == 1 and + not isinstance(pspace(expr).domain, ConditionalDomain)): + rv = tuple(condsymbols)[0] + + results = solveset(condition, rv) + if isinstance(results, Intersection) and S.Reals in results.args: + results = list(results.args[1]) + + sums = 0 + for res in results: + temp = expr.subs(rv, res) + if temp == True: + return True + if temp != False: + # XXX: This seems nonsensical but preserves existing behaviour + # after the change that Relational is no longer a subclass of + # Expr. Here expr is sometimes Relational and sometimes Expr + # but we are trying to add them with +=. This needs to be + # fixed somehow. + if sums == 0 and isinstance(expr, Relational): + sums = expr.subs(rv, res) + else: + sums += expr.subs(rv, res) + if sums == 0: + return False + return sums + + # Get full probability space of both the expression and the condition + fullspace = pspace(Tuple(expr, condition)) + # Build new space given the condition + space = fullspace.conditional_space(condition, **kwargs) + # Dictionary to swap out RandomSymbols in expr with new RandomSymbols + # That point to the new conditional space + swapdict = rs_swap(fullspace.values, space.values) + # Swap random variables in the expression + expr = expr.xreplace(swapdict) + return expr + + +def expectation(expr, condition=None, numsamples=None, evaluate=True, **kwargs): + """ + Returns the expected value of a random expression. + + Parameters + ========== + + expr : Expr containing RandomSymbols + The expression of which you want to compute the expectation value + given : Expr containing RandomSymbols + A conditional expression. E(X, X>0) is expectation of X given X > 0 + numsamples : int + Enables sampling and approximates the expectation with this many samples + evalf : Bool (defaults to True) + If sampling return a number rather than a complex expression + evaluate : Bool (defaults to True) + In case of continuous systems return unevaluated integral + + Examples + ======== + + >>> from sympy.stats import E, Die + >>> X = Die('X', 6) + >>> E(X) + 7/2 + >>> E(2*X + 1) + 8 + + >>> E(X, X > 3) # Expectation of X given that it is above 3 + 5 + """ + + if not is_random(expr): # expr isn't random? + return expr + kwargs['numsamples'] = numsamples + from sympy.stats.symbolic_probability import Expectation + if evaluate: + return Expectation(expr, condition).doit(**kwargs) + return Expectation(expr, condition) + + +def probability(condition, given_condition=None, numsamples=None, + evaluate=True, **kwargs): + """ + Probability that a condition is true, optionally given a second condition. + + Parameters + ========== + + condition : Combination of Relationals containing RandomSymbols + The condition of which you want to compute the probability + given_condition : Combination of Relationals containing RandomSymbols + A conditional expression. P(X > 1, X > 0) is expectation of X > 1 + given X > 0 + numsamples : int + Enables sampling and approximates the probability with this many samples + evaluate : Bool (defaults to True) + In case of continuous systems return unevaluated integral + + Examples + ======== + + >>> from sympy.stats import P, Die + >>> from sympy import Eq + >>> X, Y = Die('X', 6), Die('Y', 6) + >>> P(X > 3) + 1/2 + >>> P(Eq(X, 5), X > 2) # Probability that X == 5 given that X > 2 + 1/4 + >>> P(X > Y) + 5/12 + """ + + kwargs['numsamples'] = numsamples + from sympy.stats.symbolic_probability import Probability + if evaluate: + return Probability(condition, given_condition).doit(**kwargs) + return Probability(condition, given_condition) + + +class Density(Basic): + expr = property(lambda self: self.args[0]) + + def __new__(cls, expr, condition = None): + expr = _sympify(expr) + if condition is None: + obj = Basic.__new__(cls, expr) + else: + condition = _sympify(condition) + obj = Basic.__new__(cls, expr, condition) + return obj + + @property + def condition(self): + if len(self.args) > 1: + return self.args[1] + else: + return None + + def doit(self, evaluate=True, **kwargs): + from sympy.stats.random_matrix import RandomMatrixPSpace + from sympy.stats.joint_rv import JointPSpace + from sympy.stats.matrix_distributions import MatrixPSpace + from sympy.stats.compound_rv import CompoundPSpace + from sympy.stats.frv import SingleFiniteDistribution + expr, condition = self.expr, self.condition + + if isinstance(expr, SingleFiniteDistribution): + return expr.dict + if condition is not None: + # Recompute on new conditional expr + expr = given(expr, condition, **kwargs) + if not random_symbols(expr): + return Lambda(x, DiracDelta(x - expr)) + if isinstance(expr, RandomSymbol): + if isinstance(expr.pspace, (SinglePSpace, JointPSpace, MatrixPSpace)) and \ + hasattr(expr.pspace, 'distribution'): + return expr.pspace.distribution + elif isinstance(expr.pspace, RandomMatrixPSpace): + return expr.pspace.model + if isinstance(pspace(expr), CompoundPSpace): + kwargs['compound_evaluate'] = evaluate + result = pspace(expr).compute_density(expr, **kwargs) + + if evaluate and hasattr(result, 'doit'): + return result.doit() + else: + return result + + +def density(expr, condition=None, evaluate=True, numsamples=None, **kwargs): + """ + Probability density of a random expression, optionally given a second + condition. + + Explanation + =========== + + This density will take on different forms for different types of + probability spaces. Discrete variables produce Dicts. Continuous + variables produce Lambdas. + + Parameters + ========== + + expr : Expr containing RandomSymbols + The expression of which you want to compute the density value + condition : Relational containing RandomSymbols + A conditional expression. density(X > 1, X > 0) is density of X > 1 + given X > 0 + numsamples : int + Enables sampling and approximates the density with this many samples + + Examples + ======== + + >>> from sympy.stats import density, Die, Normal + >>> from sympy import Symbol + + >>> x = Symbol('x') + >>> D = Die('D', 6) + >>> X = Normal(x, 0, 1) + + >>> density(D).dict + {1: 1/6, 2: 1/6, 3: 1/6, 4: 1/6, 5: 1/6, 6: 1/6} + >>> density(2*D).dict + {2: 1/6, 4: 1/6, 6: 1/6, 8: 1/6, 10: 1/6, 12: 1/6} + >>> density(X)(x) + sqrt(2)*exp(-x**2/2)/(2*sqrt(pi)) + """ + + if numsamples: + return sampling_density(expr, condition, numsamples=numsamples, + **kwargs) + + return Density(expr, condition).doit(evaluate=evaluate, **kwargs) + + +def cdf(expr, condition=None, evaluate=True, **kwargs): + """ + Cumulative Distribution Function of a random expression. + + optionally given a second condition. + + Explanation + =========== + + This density will take on different forms for different types of + probability spaces. + Discrete variables produce Dicts. + Continuous variables produce Lambdas. + + Examples + ======== + + >>> from sympy.stats import density, Die, Normal, cdf + + >>> D = Die('D', 6) + >>> X = Normal('X', 0, 1) + + >>> density(D).dict + {1: 1/6, 2: 1/6, 3: 1/6, 4: 1/6, 5: 1/6, 6: 1/6} + >>> cdf(D) + {1: 1/6, 2: 1/3, 3: 1/2, 4: 2/3, 5: 5/6, 6: 1} + >>> cdf(3*D, D > 2) + {9: 1/4, 12: 1/2, 15: 3/4, 18: 1} + + >>> cdf(X) + Lambda(_z, erf(sqrt(2)*_z/2)/2 + 1/2) + """ + if condition is not None: # If there is a condition + # Recompute on new conditional expr + return cdf(given(expr, condition, **kwargs), **kwargs) + + # Otherwise pass work off to the ProbabilitySpace + result = pspace(expr).compute_cdf(expr, **kwargs) + + if evaluate and hasattr(result, 'doit'): + return result.doit() + else: + return result + + +def characteristic_function(expr, condition=None, evaluate=True, **kwargs): + """ + Characteristic function of a random expression, optionally given a second condition. + + Returns a Lambda. + + Examples + ======== + + >>> from sympy.stats import Normal, DiscreteUniform, Poisson, characteristic_function + + >>> X = Normal('X', 0, 1) + >>> characteristic_function(X) + Lambda(_t, exp(-_t**2/2)) + + >>> Y = DiscreteUniform('Y', [1, 2, 7]) + >>> characteristic_function(Y) + Lambda(_t, exp(7*_t*I)/3 + exp(2*_t*I)/3 + exp(_t*I)/3) + + >>> Z = Poisson('Z', 2) + >>> characteristic_function(Z) + Lambda(_t, exp(2*exp(_t*I) - 2)) + """ + if condition is not None: + return characteristic_function(given(expr, condition, **kwargs), **kwargs) + + result = pspace(expr).compute_characteristic_function(expr, **kwargs) + + if evaluate and hasattr(result, 'doit'): + return result.doit() + else: + return result + +def moment_generating_function(expr, condition=None, evaluate=True, **kwargs): + if condition is not None: + return moment_generating_function(given(expr, condition, **kwargs), **kwargs) + + result = pspace(expr).compute_moment_generating_function(expr, **kwargs) + + if evaluate and hasattr(result, 'doit'): + return result.doit() + else: + return result + +def where(condition, given_condition=None, **kwargs): + """ + Returns the domain where a condition is True. + + Examples + ======== + + >>> from sympy.stats import where, Die, Normal + >>> from sympy import And + + >>> D1, D2 = Die('a', 6), Die('b', 6) + >>> a, b = D1.symbol, D2.symbol + >>> X = Normal('x', 0, 1) + + >>> where(X**2<1) + Domain: (-1 < x) & (x < 1) + + >>> where(X**2<1).set + Interval.open(-1, 1) + + >>> where(And(D1<=D2, D2<3)) + Domain: (Eq(a, 1) & Eq(b, 1)) | (Eq(a, 1) & Eq(b, 2)) | (Eq(a, 2) & Eq(b, 2)) + """ + if given_condition is not None: # If there is a condition + # Recompute on new conditional expr + return where(given(condition, given_condition, **kwargs), **kwargs) + + # Otherwise pass work off to the ProbabilitySpace + return pspace(condition).where(condition, **kwargs) + + +@doctest_depends_on(modules=('scipy',)) +def sample(expr, condition=None, size=(), library='scipy', + numsamples=1, seed=None, **kwargs): + """ + A realization of the random expression. + + Parameters + ========== + + expr : Expression of random variables + Expression from which sample is extracted + condition : Expr containing RandomSymbols + A conditional expression + size : int, tuple + Represents size of each sample in numsamples + library : str + - 'scipy' : Sample using scipy + - 'numpy' : Sample using numpy + - 'pymc' : Sample using PyMC + + Choose any of the available options to sample from as string, + by default is 'scipy' + numsamples : int + Number of samples, each with size as ``size``. + + .. deprecated:: 1.9 + + The ``numsamples`` parameter is deprecated and is only provided for + compatibility with v1.8. Use a list comprehension or an additional + dimension in ``size`` instead. See + :ref:`deprecated-sympy-stats-numsamples` for details. + + seed : + An object to be used as seed by the given external library for sampling `expr`. + Following is the list of possible types of object for the supported libraries, + + - 'scipy': int, numpy.random.RandomState, numpy.random.Generator + - 'numpy': int, numpy.random.RandomState, numpy.random.Generator + - 'pymc': int + + Optional, by default None, in which case seed settings + related to the given library will be used. + No modifications to environment's global seed settings + are done by this argument. + + Returns + ======= + + sample: float/list/numpy.ndarray + one sample or a collection of samples of the random expression. + + - sample(X) returns float/numpy.float64/numpy.int64 object. + - sample(X, size=int/tuple) returns numpy.ndarray object. + + Examples + ======== + + >>> from sympy.stats import Die, sample, Normal, Geometric + >>> X, Y, Z = Die('X', 6), Die('Y', 6), Die('Z', 6) # Finite Random Variable + >>> die_roll = sample(X + Y + Z) + >>> die_roll # doctest: +SKIP + 3 + >>> N = Normal('N', 3, 4) # Continuous Random Variable + >>> samp = sample(N) + >>> samp in N.pspace.domain.set + True + >>> samp = sample(N, N>0) + >>> samp > 0 + True + >>> samp_list = sample(N, size=4) + >>> [sam in N.pspace.domain.set for sam in samp_list] + [True, True, True, True] + >>> sample(N, size = (2,3)) # doctest: +SKIP + array([[5.42519758, 6.40207856, 4.94991743], + [1.85819627, 6.83403519, 1.9412172 ]]) + >>> G = Geometric('G', 0.5) # Discrete Random Variable + >>> samp_list = sample(G, size=3) + >>> samp_list # doctest: +SKIP + [1, 3, 2] + >>> [sam in G.pspace.domain.set for sam in samp_list] + [True, True, True] + >>> MN = Normal("MN", [3, 4], [[2, 1], [1, 2]]) # Joint Random Variable + >>> samp_list = sample(MN, size=4) + >>> samp_list # doctest: +SKIP + [array([2.85768055, 3.38954165]), + array([4.11163337, 4.3176591 ]), + array([0.79115232, 1.63232916]), + array([4.01747268, 3.96716083])] + >>> [tuple(sam) in MN.pspace.domain.set for sam in samp_list] + [True, True, True, True] + + .. versionchanged:: 1.7.0 + sample used to return an iterator containing the samples instead of value. + + .. versionchanged:: 1.9.0 + sample returns values or array of values instead of an iterator and numsamples is deprecated. + + """ + + iterator = sample_iter(expr, condition, size=size, library=library, + numsamples=numsamples, seed=seed) + + if numsamples != 1: + sympy_deprecation_warning( + f""" + The numsamples parameter to sympy.stats.sample() is deprecated. + Either use a list comprehension, like + + [sample(...) for i in range({numsamples})] + + or add a dimension to size, like + + sample(..., size={(numsamples,) + size}) + """, + deprecated_since_version="1.9", + active_deprecations_target="deprecated-sympy-stats-numsamples", + ) + return [next(iterator) for i in range(numsamples)] + + return next(iterator) + + +def quantile(expr, evaluate=True, **kwargs): + r""" + Return the :math:`p^{th}` order quantile of a probability distribution. + + Explanation + =========== + + Quantile is defined as the value at which the probability of the random + variable is less than or equal to the given probability. + + .. math:: + Q(p) = \inf\{x \in (-\infty, \infty) : p \le F(x)\} + + Examples + ======== + + >>> from sympy.stats import quantile, Die, Exponential + >>> from sympy import Symbol, pprint + >>> p = Symbol("p") + + >>> l = Symbol("lambda", positive=True) + >>> X = Exponential("x", l) + >>> quantile(X)(p) + -log(1 - p)/lambda + + >>> D = Die("d", 6) + >>> pprint(quantile(D)(p), use_unicode=False) + /nan for Or(p > 1, p < 0) + | + | 1 for p <= 1/6 + | + | 2 for p <= 1/3 + | + < 3 for p <= 1/2 + | + | 4 for p <= 2/3 + | + | 5 for p <= 5/6 + | + \ 6 for p <= 1 + + """ + result = pspace(expr).compute_quantile(expr, **kwargs) + + if evaluate and hasattr(result, 'doit'): + return result.doit() + else: + return result + +def sample_iter(expr, condition=None, size=(), library='scipy', + numsamples=S.Infinity, seed=None, **kwargs): + + """ + Returns an iterator of realizations from the expression given a condition. + + Parameters + ========== + + expr: Expr + Random expression to be realized + condition: Expr, optional + A conditional expression + size : int, tuple + Represents size of each sample in numsamples + numsamples: integer, optional + Length of the iterator (defaults to infinity) + seed : + An object to be used as seed by the given external library for sampling `expr`. + Following is the list of possible types of object for the supported libraries, + + - 'scipy': int, numpy.random.RandomState, numpy.random.Generator + - 'numpy': int, numpy.random.RandomState, numpy.random.Generator + - 'pymc': int + + Optional, by default None, in which case seed settings + related to the given library will be used. + No modifications to environment's global seed settings + are done by this argument. + + Examples + ======== + + >>> from sympy.stats import Normal, sample_iter + >>> X = Normal('X', 0, 1) + >>> expr = X*X + 3 + >>> iterator = sample_iter(expr, numsamples=3) # doctest: +SKIP + >>> list(iterator) # doctest: +SKIP + [12, 4, 7] + + Returns + ======= + + sample_iter: iterator object + iterator object containing the sample/samples of given expr + + See Also + ======== + + sample + sampling_P + sampling_E + + """ + from sympy.stats.joint_rv import JointRandomSymbol + if not import_module(library): + raise ValueError("Failed to import %s" % library) + + if condition is not None: + ps = pspace(Tuple(expr, condition)) + else: + ps = pspace(expr) + + rvs = list(ps.values) + if isinstance(expr, JointRandomSymbol): + expr = expr.subs({expr: RandomSymbol(expr.symbol, expr.pspace)}) + else: + sub = {} + for arg in expr.args: + if isinstance(arg, JointRandomSymbol): + sub[arg] = RandomSymbol(arg.symbol, arg.pspace) + expr = expr.subs(sub) + + def fn_subs(*args): + return expr.subs(dict(zip(rvs, args))) + + def given_fn_subs(*args): + if condition is not None: + return condition.subs(dict(zip(rvs, args))) + return False + + if library in ('pymc', 'pymc3'): + # Currently unable to lambdify in pymc + # TODO : Remove when lambdify accepts 'pymc' as module + fn = lambdify(rvs, expr, **kwargs) + else: + fn = lambdify(rvs, expr, modules=library, **kwargs) + + + if condition is not None: + given_fn = lambdify(rvs, condition, **kwargs) + + def return_generator_infinite(): + count = 0 + _size = (1,)+((size,) if isinstance(size, int) else size) + while count < numsamples: + d = ps.sample(size=_size, library=library, seed=seed) # a dictionary that maps RVs to values + args = [d[rv][0] for rv in rvs] + + if condition is not None: # Check that these values satisfy the condition + # TODO: Replace the try-except block with only given_fn(*args) + # once lambdify works with unevaluated SymPy objects. + try: + gd = given_fn(*args) + except (NameError, TypeError): + gd = given_fn_subs(*args) + if gd != True and gd != False: + raise ValueError( + "Conditions must not contain free symbols") + if not gd: # If the values don't satisfy then try again + continue + + yield fn(*args) + count += 1 + + def return_generator_finite(): + faulty = True + while faulty: + d = ps.sample(size=(numsamples,) + ((size,) if isinstance(size, int) else size), + library=library, seed=seed) # a dictionary that maps RVs to values + + faulty = False + count = 0 + while count < numsamples and not faulty: + args = [d[rv][count] for rv in rvs] + if condition is not None: # Check that these values satisfy the condition + # TODO: Replace the try-except block with only given_fn(*args) + # once lambdify works with unevaluated SymPy objects. + try: + gd = given_fn(*args) + except (NameError, TypeError): + gd = given_fn_subs(*args) + if gd != True and gd != False: + raise ValueError( + "Conditions must not contain free symbols") + if not gd: # If the values don't satisfy then try again + faulty = True + + count += 1 + + count = 0 + while count < numsamples: + args = [d[rv][count] for rv in rvs] + # TODO: Replace the try-except block with only fn(*args) + # once lambdify works with unevaluated SymPy objects. + try: + yield fn(*args) + except (NameError, TypeError): + yield fn_subs(*args) + count += 1 + + if numsamples is S.Infinity: + return return_generator_infinite() + + return return_generator_finite() + +def sample_iter_lambdify(expr, condition=None, size=(), + numsamples=S.Infinity, seed=None, **kwargs): + + return sample_iter(expr, condition=condition, size=size, + numsamples=numsamples, seed=seed, **kwargs) + +def sample_iter_subs(expr, condition=None, size=(), + numsamples=S.Infinity, seed=None, **kwargs): + + return sample_iter(expr, condition=condition, size=size, + numsamples=numsamples, seed=seed, **kwargs) + + +def sampling_P(condition, given_condition=None, library='scipy', numsamples=1, + evalf=True, seed=None, **kwargs): + """ + Sampling version of P. + + See Also + ======== + + P + sampling_E + sampling_density + + """ + + count_true = 0 + count_false = 0 + samples = sample_iter(condition, given_condition, library=library, + numsamples=numsamples, seed=seed, **kwargs) + + for sample in samples: + if sample: + count_true += 1 + else: + count_false += 1 + + result = S(count_true) / numsamples + if evalf: + return result.evalf() + else: + return result + + +def sampling_E(expr, given_condition=None, library='scipy', numsamples=1, + evalf=True, seed=None, **kwargs): + """ + Sampling version of E. + + See Also + ======== + + P + sampling_P + sampling_density + """ + samples = list(sample_iter(expr, given_condition, library=library, + numsamples=numsamples, seed=seed, **kwargs)) + result = Add(*samples) / numsamples + + if evalf: + return result.evalf() + else: + return result + +def sampling_density(expr, given_condition=None, library='scipy', + numsamples=1, seed=None, **kwargs): + """ + Sampling version of density. + + See Also + ======== + density + sampling_P + sampling_E + """ + + results = {} + for result in sample_iter(expr, given_condition, library=library, + numsamples=numsamples, seed=seed, **kwargs): + results[result] = results.get(result, 0) + 1 + + return results + + +def dependent(a, b): + """ + Dependence of two random expressions. + + Two expressions are independent if knowledge of one does not change + computations on the other. + + Examples + ======== + + >>> from sympy.stats import Normal, dependent, given + >>> from sympy import Tuple, Eq + + >>> X, Y = Normal('X', 0, 1), Normal('Y', 0, 1) + >>> dependent(X, Y) + False + >>> dependent(2*X + Y, -Y) + True + >>> X, Y = given(Tuple(X, Y), Eq(X + Y, 3)) + >>> dependent(X, Y) + True + + See Also + ======== + + independent + """ + if pspace_independent(a, b): + return False + + z = Symbol('z', real=True) + # Dependent if density is unchanged when one is given information about + # the other + return (density(a, Eq(b, z)) != density(a) or + density(b, Eq(a, z)) != density(b)) + + +def independent(a, b): + """ + Independence of two random expressions. + + Two expressions are independent if knowledge of one does not change + computations on the other. + + Examples + ======== + + >>> from sympy.stats import Normal, independent, given + >>> from sympy import Tuple, Eq + + >>> X, Y = Normal('X', 0, 1), Normal('Y', 0, 1) + >>> independent(X, Y) + True + >>> independent(2*X + Y, -Y) + False + >>> X, Y = given(Tuple(X, Y), Eq(X + Y, 3)) + >>> independent(X, Y) + False + + See Also + ======== + + dependent + """ + return not dependent(a, b) + + +def pspace_independent(a, b): + """ + Tests for independence between a and b by checking if their PSpaces have + overlapping symbols. This is a sufficient but not necessary condition for + independence and is intended to be used internally. + + Notes + ===== + + pspace_independent(a, b) implies independent(a, b) + independent(a, b) does not imply pspace_independent(a, b) + """ + a_symbols = set(pspace(b).symbols) + b_symbols = set(pspace(a).symbols) + + if len(set(random_symbols(a)).intersection(random_symbols(b))) != 0: + return False + + if len(a_symbols.intersection(b_symbols)) == 0: + return True + return None + + +def rv_subs(expr, symbols=None): + """ + Given a random expression replace all random variables with their symbols. + + If symbols keyword is given restrict the swap to only the symbols listed. + """ + if symbols is None: + symbols = random_symbols(expr) + if not symbols: + return expr + swapdict = {rv: rv.symbol for rv in symbols} + return expr.subs(swapdict) + + +class NamedArgsMixin: + _argnames: tuple[str, ...] = () + + def __getattr__(self, attr): + try: + return self.args[self._argnames.index(attr)] + except ValueError: + raise AttributeError("'%s' object has no attribute '%s'" % ( + type(self).__name__, attr)) + + +class Distribution(Basic): + + def sample(self, size=(), library='scipy', seed=None): + """ A random realization from the distribution """ + + module = import_module(library) + if library in {'scipy', 'numpy', 'pymc3', 'pymc'} and module is None: + raise ValueError("Failed to import %s" % library) + + if library == 'scipy': + # scipy does not require map as it can handle using custom distributions. + # However, we will still use a map where we can. + + # TODO: do this for drv.py and frv.py if necessary. + # TODO: add more distributions here if there are more + # See links below referring to sections beginning with "A common parametrization..." + # I will remove all these comments if everything is ok. + + from sympy.stats.sampling.sample_scipy import do_sample_scipy + import numpy + if seed is None or isinstance(seed, int): + rand_state = numpy.random.default_rng(seed=seed) + else: + rand_state = seed + samps = do_sample_scipy(self, size, rand_state) + + elif library == 'numpy': + from sympy.stats.sampling.sample_numpy import do_sample_numpy + import numpy + if seed is None or isinstance(seed, int): + rand_state = numpy.random.default_rng(seed=seed) + else: + rand_state = seed + _size = None if size == () else size + samps = do_sample_numpy(self, _size, rand_state) + elif library in ('pymc', 'pymc3'): + from sympy.stats.sampling.sample_pymc import do_sample_pymc + import logging + logging.getLogger("pymc").setLevel(logging.ERROR) + try: + import pymc + except ImportError: + import pymc3 as pymc + + with pymc.Model(): + if do_sample_pymc(self): + samps = pymc.sample(draws=prod(size), chains=1, compute_convergence_checks=False, + progressbar=False, random_seed=seed, return_inferencedata=False)[:]['X'] + samps = samps.reshape(size) + else: + samps = None + + else: + raise NotImplementedError("Sampling from %s is not supported yet." + % str(library)) + + if samps is not None: + return samps + raise NotImplementedError( + "Sampling for %s is not currently implemented from %s" + % (self, library)) + + +def _value_check(condition, message): + """ + Raise a ValueError with message if condition is False, else + return True if all conditions were True, else False. + + Examples + ======== + + >>> from sympy.stats.rv import _value_check + >>> from sympy.abc import a, b, c + >>> from sympy import And, Dummy + + >>> _value_check(2 < 3, '') + True + + Here, the condition is not False, but it does not evaluate to True + so False is returned (but no error is raised). So checking if the + return value is True or False will tell you if all conditions were + evaluated. + + >>> _value_check(a < b, '') + False + + In this case the condition is False so an error is raised: + + >>> r = Dummy(real=True) + >>> _value_check(r < r - 1, 'condition is not true') + Traceback (most recent call last): + ... + ValueError: condition is not true + + If no condition of many conditions must be False, they can be + checked by passing them as an iterable: + + >>> _value_check((a < 0, b < 0, c < 0), '') + False + + The iterable can be a generator, too: + + >>> _value_check((i < 0 for i in (a, b, c)), '') + False + + The following are equivalent to the above but do not pass + an iterable: + + >>> all(_value_check(i < 0, '') for i in (a, b, c)) + False + >>> _value_check(And(a < 0, b < 0, c < 0), '') + False + """ + if not iterable(condition): + condition = [condition] + truth = fuzzy_and(condition) + if truth == False: + raise ValueError(message) + return truth == True + +def _symbol_converter(sym): + """ + Casts the parameter to Symbol if it is 'str' + otherwise no operation is performed on it. + + Parameters + ========== + + sym + The parameter to be converted. + + Returns + ======= + + Symbol + the parameter converted to Symbol. + + Raises + ====== + + TypeError + If the parameter is not an instance of both str and + Symbol. + + Examples + ======== + + >>> from sympy import Symbol + >>> from sympy.stats.rv import _symbol_converter + >>> s = _symbol_converter('s') + >>> isinstance(s, Symbol) + True + >>> _symbol_converter(1) + Traceback (most recent call last): + ... + TypeError: 1 is neither a Symbol nor a string + >>> r = Symbol('r') + >>> isinstance(r, Symbol) + True + """ + if isinstance(sym, str): + sym = Symbol(sym) + if not isinstance(sym, Symbol): + raise TypeError("%s is neither a Symbol nor a string"%(sym)) + return sym + +def sample_stochastic_process(process): + """ + This function is used to sample from stochastic process. + + Parameters + ========== + + process: StochasticProcess + Process used to extract the samples. It must be an instance of + StochasticProcess + + Examples + ======== + + >>> from sympy.stats import sample_stochastic_process, DiscreteMarkovChain + >>> from sympy import Matrix + >>> T = Matrix([[0.5, 0.2, 0.3],[0.2, 0.5, 0.3],[0.2, 0.3, 0.5]]) + >>> Y = DiscreteMarkovChain("Y", [0, 1, 2], T) + >>> next(sample_stochastic_process(Y)) in Y.state_space + True + >>> next(sample_stochastic_process(Y)) # doctest: +SKIP + 0 + >>> next(sample_stochastic_process(Y)) # doctest: +SKIP + 2 + + Returns + ======= + + sample: iterator object + iterator object containing the sample of given process + + """ + from sympy.stats.stochastic_process_types import StochasticProcess + if not isinstance(process, StochasticProcess): + raise ValueError("Process must be an instance of Stochastic Process") + return process.sample() diff --git a/lib/python3.12/site-packages/sympy/stats/rv_interface.py b/lib/python3.12/site-packages/sympy/stats/rv_interface.py new file mode 100644 index 0000000000000000000000000000000000000000..4a4a6715d7d54b9553afac132ac3f577f095a137 --- /dev/null +++ b/lib/python3.12/site-packages/sympy/stats/rv_interface.py @@ -0,0 +1,519 @@ +from sympy.sets import FiniteSet +from sympy.core.numbers import Rational +from sympy.core.relational import Eq +from sympy.core.symbol import Dummy +from sympy.functions.combinatorial.factorials import FallingFactorial +from sympy.functions.elementary.exponential import (exp, log) +from sympy.functions.elementary.miscellaneous import sqrt +from sympy.functions.elementary.piecewise import piecewise_fold +from sympy.integrals.integrals import Integral +from sympy.solvers.solveset import solveset +from .rv import (probability, expectation, density, where, given, pspace, cdf, PSpace, + characteristic_function, sample, sample_iter, random_symbols, independent, dependent, + sampling_density, moment_generating_function, quantile, is_random, + sample_stochastic_process) + + +__all__ = ['P', 'E', 'H', 'density', 'where', 'given', 'sample', 'cdf', + 'characteristic_function', 'pspace', 'sample_iter', 'variance', 'std', + 'skewness', 'kurtosis', 'covariance', 'dependent', 'entropy', 'median', + 'independent', 'random_symbols', 'correlation', 'factorial_moment', + 'moment', 'cmoment', 'sampling_density', 'moment_generating_function', + 'smoment', 'quantile', 'sample_stochastic_process'] + + + +def moment(X, n, c=0, condition=None, *, evaluate=True, **kwargs): + """ + Return the nth moment of a random expression about c. + + .. math:: + moment(X, c, n) = E((X-c)^{n}) + + Default value of c is 0. + + Examples + ======== + + >>> from sympy.stats import Die, moment, E + >>> X = Die('X', 6) + >>> moment(X, 1, 6) + -5/2 + >>> moment(X, 2) + 91/6 + >>> moment(X, 1) == E(X) + True + """ + from sympy.stats.symbolic_probability import Moment + if evaluate: + return Moment(X, n, c, condition).doit() + return Moment(X, n, c, condition).rewrite(Integral) + + +def variance(X, condition=None, **kwargs): + """ + Variance of a random expression. + + .. math:: + variance(X) = E((X-E(X))^{2}) + + Examples + ======== + + >>> from sympy.stats import Die, Bernoulli, variance + >>> from sympy import simplify, Symbol + + >>> X = Die('X', 6) + >>> p = Symbol('p') + >>> B = Bernoulli('B', p, 1, 0) + + >>> variance(2*X) + 35/3 + + >>> simplify(variance(B)) + p*(1 - p) + """ + if is_random(X) and pspace(X) == PSpace(): + from sympy.stats.symbolic_probability import Variance + return Variance(X, condition) + + return cmoment(X, 2, condition, **kwargs) + + +def standard_deviation(X, condition=None, **kwargs): + r""" + Standard Deviation of a random expression + + .. math:: + std(X) = \sqrt(E((X-E(X))^{2})) + + Examples + ======== + + >>> from sympy.stats import Bernoulli, std + >>> from sympy import Symbol, simplify + + >>> p = Symbol('p') + >>> B = Bernoulli('B', p, 1, 0) + + >>> simplify(std(B)) + sqrt(p*(1 - p)) + """ + return sqrt(variance(X, condition, **kwargs)) +std = standard_deviation + +def entropy(expr, condition=None, **kwargs): + """ + Calculuates entropy of a probability distribution. + + Parameters + ========== + + expression : the random expression whose entropy is to be calculated + condition : optional, to specify conditions on random expression + b: base of the logarithm, optional + By default, it is taken as Euler's number + + Returns + ======= + + result : Entropy of the expression, a constant + + Examples + ======== + + >>> from sympy.stats import Normal, Die, entropy + >>> X = Normal('X', 0, 1) + >>> entropy(X) + log(2)/2 + 1/2 + log(pi)/2 + + >>> D = Die('D', 4) + >>> entropy(D) + log(4) + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Entropy_%28information_theory%29 + .. [2] https://www.crmarsh.com/static/pdf/Charles_Marsh_Continuous_Entropy.pdf + .. [3] https://kconrad.math.uconn.edu/blurbs/analysis/entropypost.pdf + """ + pdf = density(expr, condition, **kwargs) + base = kwargs.get('b', exp(1)) + if isinstance(pdf, dict): + return sum(-prob*log(prob, base) for prob in pdf.values()) + return expectation(-log(pdf(expr), base)) + +def covariance(X, Y, condition=None, **kwargs): + """ + Covariance of two random expressions. + + Explanation + =========== + + The expectation that the two variables will rise and fall together + + .. math:: + covariance(X,Y) = E((X-E(X)) (Y-E(Y))) + + Examples + ======== + + >>> from sympy.stats import Exponential, covariance + >>> from sympy import Symbol + + >>> rate = Symbol('lambda', positive=True, real=True) + >>> X = Exponential('X', rate) + >>> Y = Exponential('Y', rate) + + >>> covariance(X, X) + lambda**(-2) + >>> covariance(X, Y) + 0 + >>> covariance(X, Y + rate*X) + 1/lambda + """ + if (is_random(X) and pspace(X) == PSpace()) or (is_random(Y) and pspace(Y) == PSpace()): + from sympy.stats.symbolic_probability import Covariance + return Covariance(X, Y, condition) + + return expectation( + (X - expectation(X, condition, **kwargs)) * + (Y - expectation(Y, condition, **kwargs)), + condition, **kwargs) + + +def correlation(X, Y, condition=None, **kwargs): + r""" + Correlation of two random expressions, also known as correlation + coefficient or Pearson's correlation. + + Explanation + =========== + + The normalized expectation that the two variables will rise + and fall together + + .. math:: + correlation(X,Y) = E((X-E(X))(Y-E(Y)) / (\sigma_x \sigma_y)) + + Examples + ======== + + >>> from sympy.stats import Exponential, correlation + >>> from sympy import Symbol + + >>> rate = Symbol('lambda', positive=True, real=True) + >>> X = Exponential('X', rate) + >>> Y = Exponential('Y', rate) + + >>> correlation(X, X) + 1 + >>> correlation(X, Y) + 0 + >>> correlation(X, Y + rate*X) + 1/sqrt(1 + lambda**(-2)) + """ + return covariance(X, Y, condition, **kwargs)/(std(X, condition, **kwargs) + * std(Y, condition, **kwargs)) + + +def cmoment(X, n, condition=None, *, evaluate=True, **kwargs): + """ + Return the nth central moment of a random expression about its mean. + + .. math:: + cmoment(X, n) = E((X - E(X))^{n}) + + Examples + ======== + + >>> from sympy.stats import Die, cmoment, variance + >>> X = Die('X', 6) + >>> cmoment(X, 3) + 0 + >>> cmoment(X, 2) + 35/12 + >>> cmoment(X, 2) == variance(X) + True + """ + from sympy.stats.symbolic_probability import CentralMoment + if evaluate: + return CentralMoment(X, n, condition).doit() + return CentralMoment(X, n, condition).rewrite(Integral) + + +def smoment(X, n, condition=None, **kwargs): + r""" + Return the nth Standardized moment of a random expression. + + .. math:: + smoment(X, n) = E(((X - \mu)/\sigma_X)^{n}) + + Examples + ======== + + >>> from sympy.stats import skewness, Exponential, smoment + >>> from sympy import Symbol + >>> rate = Symbol('lambda', positive=True, real=True) + >>> Y = Exponential('Y', rate) + >>> smoment(Y, 4) + 9 + >>> smoment(Y, 4) == smoment(3*Y, 4) + True + >>> smoment(Y, 3) == skewness(Y) + True + """ + sigma = std(X, condition, **kwargs) + return (1/sigma)**n*cmoment(X, n, condition, **kwargs) + +def skewness(X, condition=None, **kwargs): + r""" + Measure of the asymmetry of the probability distribution. + + Explanation + =========== + + Positive skew indicates that most of the values lie to the right of + the mean. + + .. math:: + skewness(X) = E(((X - E(X))/\sigma_X)^{3}) + + Parameters + ========== + + condition : Expr containing RandomSymbols + A conditional expression. skewness(X, X>0) is skewness of X given X > 0 + + Examples + ======== + + >>> from sympy.stats import skewness, Exponential, Normal + >>> from sympy import Symbol + >>> X = Normal('X', 0, 1) + >>> skewness(X) + 0 + >>> skewness(X, X > 0) # find skewness given X > 0 + (-sqrt(2)/sqrt(pi) + 4*sqrt(2)/pi**(3/2))/(1 - 2/pi)**(3/2) + + >>> rate = Symbol('lambda', positive=True, real=True) + >>> Y = Exponential('Y', rate) + >>> skewness(Y) + 2 + """ + return smoment(X, 3, condition=condition, **kwargs) + +def kurtosis(X, condition=None, **kwargs): + r""" + Characterizes the tails/outliers of a probability distribution. + + Explanation + =========== + + Kurtosis of any univariate normal distribution is 3. Kurtosis less than + 3 means that the distribution produces fewer and less extreme outliers + than the normal distribution. + + .. math:: + kurtosis(X) = E(((X - E(X))/\sigma_X)^{4}) + + Parameters + ========== + + condition : Expr containing RandomSymbols + A conditional expression. kurtosis(X, X>0) is kurtosis of X given X > 0 + + Examples + ======== + + >>> from sympy.stats import kurtosis, Exponential, Normal + >>> from sympy import Symbol + >>> X = Normal('X', 0, 1) + >>> kurtosis(X) + 3 + >>> kurtosis(X, X > 0) # find kurtosis given X > 0 + (-4/pi - 12/pi**2 + 3)/(1 - 2/pi)**2 + + >>> rate = Symbol('lamda', positive=True, real=True) + >>> Y = Exponential('Y', rate) + >>> kurtosis(Y) + 9 + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Kurtosis + .. [2] https://mathworld.wolfram.com/Kurtosis.html + """ + return smoment(X, 4, condition=condition, **kwargs) + + +def factorial_moment(X, n, condition=None, **kwargs): + """ + The factorial moment is a mathematical quantity defined as the expectation + or average of the falling factorial of a random variable. + + .. math:: + factorial-moment(X, n) = E(X(X - 1)(X - 2)...(X - n + 1)) + + Parameters + ========== + + n: A natural number, n-th factorial moment. + + condition : Expr containing RandomSymbols + A conditional expression. + + Examples + ======== + + >>> from sympy.stats import factorial_moment, Poisson, Binomial + >>> from sympy import Symbol, S + >>> lamda = Symbol('lamda') + >>> X = Poisson('X', lamda) + >>> factorial_moment(X, 2) + lamda**2 + >>> Y = Binomial('Y', 2, S.Half) + >>> factorial_moment(Y, 2) + 1/2 + >>> factorial_moment(Y, 2, Y > 1) # find factorial moment for Y > 1 + 2 + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Factorial_moment + .. [2] https://mathworld.wolfram.com/FactorialMoment.html + """ + return expectation(FallingFactorial(X, n), condition=condition, **kwargs) + +def median(X, evaluate=True, **kwargs): + r""" + Calculuates the median of the probability distribution. + + Explanation + =========== + + Mathematically, median of Probability distribution is defined as all those + values of `m` for which the following condition is satisfied + + .. math:: + P(X\leq m) \geq \frac{1}{2} \text{ and} \text{ } P(X\geq m)\geq \frac{1}{2} + + Parameters + ========== + + X: The random expression whose median is to be calculated. + + Returns + ======= + + The FiniteSet or an Interval which contains the median of the + random expression. + + Examples + ======== + + >>> from sympy.stats import Normal, Die, median + >>> N = Normal('N', 3, 1) + >>> median(N) + {3} + >>> D = Die('D') + >>> median(D) + {3, 4} + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Median#Probability_distributions + + """ + if not is_random(X): + return X + + from sympy.stats.crv import ContinuousPSpace + from sympy.stats.drv import DiscretePSpace + from sympy.stats.frv import FinitePSpace + + if isinstance(pspace(X), FinitePSpace): + cdf = pspace(X).compute_cdf(X) + result = [] + for key, value in cdf.items(): + if value>= Rational(1, 2) and (1 - value) + \ + pspace(X).probability(Eq(X, key)) >= Rational(1, 2): + result.append(key) + return FiniteSet(*result) + if isinstance(pspace(X), (ContinuousPSpace, DiscretePSpace)): + cdf = pspace(X).compute_cdf(X) + x = Dummy('x') + result = solveset(piecewise_fold(cdf(x) - Rational(1, 2)), x, pspace(X).set) + return result + raise NotImplementedError("The median of %s is not implemented."%str(pspace(X))) + + +def coskewness(X, Y, Z, condition=None, **kwargs): + r""" + Calculates the co-skewness of three random variables. + + Explanation + =========== + + Mathematically Coskewness is defined as + + .. math:: + coskewness(X,Y,Z)=\frac{E[(X-E[X]) * (Y-E[Y]) * (Z-E[Z])]} {\sigma_{X}\sigma_{Y}\sigma_{Z}} + + Parameters + ========== + + X : RandomSymbol + Random Variable used to calculate coskewness + Y : RandomSymbol + Random Variable used to calculate coskewness + Z : RandomSymbol + Random Variable used to calculate coskewness + condition : Expr containing RandomSymbols + A conditional expression + + Examples + ======== + + >>> from sympy.stats import coskewness, Exponential, skewness + >>> from sympy import symbols + >>> p = symbols('p', positive=True) + >>> X = Exponential('X', p) + >>> Y = Exponential('Y', 2*p) + >>> coskewness(X, Y, Y) + 0 + >>> coskewness(X, Y + X, Y + 2*X) + 16*sqrt(85)/85 + >>> coskewness(X + 2*Y, Y + X, Y + 2*X, X > 3) + 9*sqrt(170)/85 + >>> coskewness(Y, Y, Y) == skewness(Y) + True + >>> coskewness(X, Y + p*X, Y + 2*p*X) + 4/(sqrt(1 + 1/(4*p**2))*sqrt(4 + 1/(4*p**2))) + + Returns + ======= + + coskewness : The coskewness of the three random variables + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Coskewness + + """ + num = expectation((X - expectation(X, condition, **kwargs)) \ + * (Y - expectation(Y, condition, **kwargs)) \ + * (Z - expectation(Z, condition, **kwargs)), condition, **kwargs) + den = std(X, condition, **kwargs) * std(Y, condition, **kwargs) \ + * std(Z, condition, **kwargs) + return num/den + + +P = probability +E = expectation +H = entropy diff --git a/lib/python3.12/site-packages/sympy/stats/stochastic_process.py b/lib/python3.12/site-packages/sympy/stats/stochastic_process.py new file mode 100644 index 0000000000000000000000000000000000000000..bfb0e759c66be892ae38ddda004dfe928f683fee --- /dev/null +++ b/lib/python3.12/site-packages/sympy/stats/stochastic_process.py @@ -0,0 +1,66 @@ +from sympy.core.basic import Basic +from sympy.stats.joint_rv import ProductPSpace +from sympy.stats.rv import ProductDomain, _symbol_converter, Distribution + + +class StochasticPSpace(ProductPSpace): + """ + Represents probability space of stochastic processes + and their random variables. Contains mechanics to do + computations for queries of stochastic processes. + + Explanation + =========== + + Initialized by symbol, the specific process and + distribution(optional) if the random indexed symbols + of the process follows any specific distribution, like, + in Bernoulli Process, each random indexed symbol follows + Bernoulli distribution. For processes with memory, this + parameter should not be passed. + """ + + def __new__(cls, sym, process, distribution=None): + sym = _symbol_converter(sym) + from sympy.stats.stochastic_process_types import StochasticProcess + if not isinstance(process, StochasticProcess): + raise TypeError("`process` must be an instance of StochasticProcess.") + if distribution is None: + distribution = Distribution() + return Basic.__new__(cls, sym, process, distribution) + + @property + def process(self): + """ + The associated stochastic process. + """ + return self.args[1] + + @property + def domain(self): + return ProductDomain(self.process.index_set, + self.process.state_space) + + @property + def symbol(self): + return self.args[0] + + @property + def distribution(self): + return self.args[2] + + def probability(self, condition, given_condition=None, evaluate=True, **kwargs): + """ + Transfers the task of handling queries to the specific stochastic + process because every process has their own logic of handling such + queries. + """ + return self.process.probability(condition, given_condition, evaluate, **kwargs) + + def compute_expectation(self, expr, condition=None, evaluate=True, **kwargs): + """ + Transfers the task of handling queries to the specific stochastic + process because every process has their own logic of handling such + queries. + """ + return self.process.expectation(expr, condition, evaluate, **kwargs) diff --git a/lib/python3.12/site-packages/sympy/stats/stochastic_process_types.py b/lib/python3.12/site-packages/sympy/stats/stochastic_process_types.py new file mode 100644 index 0000000000000000000000000000000000000000..9bebed5d2052190a38b6ff5d1b71f98f196b53cc --- /dev/null +++ b/lib/python3.12/site-packages/sympy/stats/stochastic_process_types.py @@ -0,0 +1,2383 @@ +import random +import itertools +from typing import (Sequence as tSequence, Union as tUnion, List as tList, + Tuple as tTuple, Set as tSet) +from sympy.concrete.summations import Sum +from sympy.core.add import Add +from sympy.core.basic import Basic +from sympy.core.cache import cacheit +from sympy.core.containers import Tuple +from sympy.core.expr import Expr +from sympy.core.function import (Function, Lambda) +from sympy.core.mul import Mul +from sympy.core.intfunc import igcd +from sympy.core.numbers import (Integer, Rational, oo, pi) +from sympy.core.relational import (Eq, Ge, Gt, Le, Lt, Ne) +from sympy.core.singleton import S +from sympy.core.symbol import (Dummy, Symbol) +from sympy.functions.combinatorial.factorials import factorial +from sympy.functions.elementary.exponential import exp +from sympy.functions.elementary.integers import ceiling +from sympy.functions.elementary.miscellaneous import sqrt +from sympy.functions.elementary.piecewise import Piecewise +from sympy.functions.special.gamma_functions import gamma +from sympy.logic.boolalg import (And, Not, Or) +from sympy.matrices.exceptions import NonSquareMatrixError +from sympy.matrices.dense import (Matrix, eye, ones, zeros) +from sympy.matrices.expressions.blockmatrix import BlockMatrix +from sympy.matrices.expressions.matexpr import MatrixSymbol +from sympy.matrices.expressions.special import Identity +from sympy.matrices.immutable import ImmutableMatrix +from sympy.sets.conditionset import ConditionSet +from sympy.sets.contains import Contains +from sympy.sets.fancysets import Range +from sympy.sets.sets import (FiniteSet, Intersection, Interval, Set, Union) +from sympy.solvers.solveset import linsolve +from sympy.tensor.indexed import (Indexed, IndexedBase) +from sympy.core.relational import Relational +from sympy.logic.boolalg import Boolean +from sympy.utilities.exceptions import sympy_deprecation_warning +from sympy.utilities.iterables import strongly_connected_components +from sympy.stats.joint_rv import JointDistribution +from sympy.stats.joint_rv_types import JointDistributionHandmade +from sympy.stats.rv import (RandomIndexedSymbol, random_symbols, RandomSymbol, + _symbol_converter, _value_check, pspace, given, + dependent, is_random, sample_iter, Distribution, + Density) +from sympy.stats.stochastic_process import StochasticPSpace +from sympy.stats.symbolic_probability import Probability, Expectation +from sympy.stats.frv_types import Bernoulli, BernoulliDistribution, FiniteRV +from sympy.stats.drv_types import Poisson, PoissonDistribution +from sympy.stats.crv_types import Normal, NormalDistribution, Gamma, GammaDistribution +from sympy.core.sympify import _sympify, sympify + +EmptySet = S.EmptySet + +__all__ = [ + 'StochasticProcess', + 'DiscreteTimeStochasticProcess', + 'DiscreteMarkovChain', + 'TransitionMatrixOf', + 'StochasticStateSpaceOf', + 'GeneratorMatrixOf', + 'ContinuousMarkovChain', + 'BernoulliProcess', + 'PoissonProcess', + 'WienerProcess', + 'GammaProcess' +] + + +@is_random.register(Indexed) +def _(x): + return is_random(x.base) + +@is_random.register(RandomIndexedSymbol) # type: ignore +def _(x): + return True + +def _set_converter(itr): + """ + Helper function for converting list/tuple/set to Set. + If parameter is not an instance of list/tuple/set then + no operation is performed. + + Returns + ======= + + Set + The argument converted to Set. + + + Raises + ====== + + TypeError + If the argument is not an instance of list/tuple/set. + """ + if isinstance(itr, (list, tuple, set)): + itr = FiniteSet(*itr) + if not isinstance(itr, Set): + raise TypeError("%s is not an instance of list/tuple/set."%(itr)) + return itr + +def _state_converter(itr: tSequence) -> tUnion[Tuple, Range]: + """ + Helper function for converting list/tuple/set/Range/Tuple/FiniteSet + to tuple/Range. + """ + itr_ret: tUnion[Tuple, Range] + + if isinstance(itr, (Tuple, set, FiniteSet)): + itr_ret = Tuple(*(sympify(i) if isinstance(i, str) else i for i in itr)) + + elif isinstance(itr, (list, tuple)): + # check if states are unique + if len(set(itr)) != len(itr): + raise ValueError('The state space must have unique elements.') + itr_ret = Tuple(*(sympify(i) if isinstance(i, str) else i for i in itr)) + + elif isinstance(itr, Range): + # the only ordered set in SymPy I know of + # try to convert to tuple + try: + itr_ret = Tuple(*(sympify(i) if isinstance(i, str) else i for i in itr)) + except (TypeError, ValueError): + itr_ret = itr + + else: + raise TypeError("%s is not an instance of list/tuple/set/Range/Tuple/FiniteSet." % (itr)) + return itr_ret + +def _sym_sympify(arg): + """ + Converts an arbitrary expression to a type that can be used inside SymPy. + As generally strings are unwise to use in the expressions, + it returns the Symbol of argument if the string type argument is passed. + + Parameters + ========= + + arg: The parameter to be converted to be used in SymPy. + + Returns + ======= + + The converted parameter. + + """ + if isinstance(arg, str): + return Symbol(arg) + else: + return _sympify(arg) + +def _matrix_checks(matrix): + if not isinstance(matrix, (Matrix, MatrixSymbol, ImmutableMatrix)): + raise TypeError("Transition probabilities either should " + "be a Matrix or a MatrixSymbol.") + if matrix.shape[0] != matrix.shape[1]: + raise NonSquareMatrixError("%s is not a square matrix"%(matrix)) + if isinstance(matrix, Matrix): + matrix = ImmutableMatrix(matrix.tolist()) + return matrix + +class StochasticProcess(Basic): + """ + Base class for all the stochastic processes whether + discrete or continuous. + + Parameters + ========== + + sym: Symbol or str + state_space: Set + The state space of the stochastic process, by default S.Reals. + For discrete sets it is zero indexed. + + See Also + ======== + + DiscreteTimeStochasticProcess + """ + + index_set = S.Reals + + def __new__(cls, sym, state_space=S.Reals, **kwargs): + sym = _symbol_converter(sym) + state_space = _set_converter(state_space) + return Basic.__new__(cls, sym, state_space) + + @property + def symbol(self): + return self.args[0] + + @property + def state_space(self) -> tUnion[FiniteSet, Range]: + if not isinstance(self.args[1], (FiniteSet, Range)): + assert isinstance(self.args[1], Tuple) + return FiniteSet(*self.args[1]) + return self.args[1] + + def _deprecation_warn_distribution(self): + sympy_deprecation_warning( + """ + Calling the distribution method with a RandomIndexedSymbol + argument, like X.distribution(X(t)) is deprecated. Instead, call + distribution() with the given timestamp, like + + X.distribution(t) + """, + deprecated_since_version="1.7.1", + active_deprecations_target="deprecated-distribution-randomindexedsymbol", + stacklevel=4, + ) + + def distribution(self, key=None): + if key is None: + self._deprecation_warn_distribution() + return Distribution() + + def density(self, x): + return Density() + + def __call__(self, time): + """ + Overridden in ContinuousTimeStochasticProcess. + """ + raise NotImplementedError("Use [] for indexing discrete time stochastic process.") + + def __getitem__(self, time): + """ + Overridden in DiscreteTimeStochasticProcess. + """ + raise NotImplementedError("Use () for indexing continuous time stochastic process.") + + def probability(self, condition): + raise NotImplementedError() + + def joint_distribution(self, *args): + """ + Computes the joint distribution of the random indexed variables. + + Parameters + ========== + + args: iterable + The finite list of random indexed variables/the key of a stochastic + process whose joint distribution has to be computed. + + Returns + ======= + + JointDistribution + The joint distribution of the list of random indexed variables. + An unevaluated object is returned if it is not possible to + compute the joint distribution. + + Raises + ====== + + ValueError: When the arguments passed are not of type RandomIndexSymbol + or Number. + """ + args = list(args) + for i, arg in enumerate(args): + if S(arg).is_Number: + if self.index_set.is_subset(S.Integers): + args[i] = self.__getitem__(arg) + else: + args[i] = self.__call__(arg) + elif not isinstance(arg, RandomIndexedSymbol): + raise ValueError("Expected a RandomIndexedSymbol or " + "key not %s"%(type(arg))) + + if args[0].pspace.distribution == Distribution(): + return JointDistribution(*args) + density = Lambda(tuple(args), + expr=Mul.fromiter(arg.pspace.process.density(arg) for arg in args)) + return JointDistributionHandmade(density) + + def expectation(self, condition, given_condition): + raise NotImplementedError("Abstract method for expectation queries.") + + def sample(self): + raise NotImplementedError("Abstract method for sampling queries.") + +class DiscreteTimeStochasticProcess(StochasticProcess): + """ + Base class for all discrete stochastic processes. + """ + def __getitem__(self, time): + """ + For indexing discrete time stochastic processes. + + Returns + ======= + + RandomIndexedSymbol + """ + time = sympify(time) + if not time.is_symbol and time not in self.index_set: + raise IndexError("%s is not in the index set of %s"%(time, self.symbol)) + idx_obj = Indexed(self.symbol, time) + pspace_obj = StochasticPSpace(self.symbol, self, self.distribution(time)) + return RandomIndexedSymbol(idx_obj, pspace_obj) + +class ContinuousTimeStochasticProcess(StochasticProcess): + """ + Base class for all continuous time stochastic process. + """ + def __call__(self, time): + """ + For indexing continuous time stochastic processes. + + Returns + ======= + + RandomIndexedSymbol + """ + time = sympify(time) + if not time.is_symbol and time not in self.index_set: + raise IndexError("%s is not in the index set of %s"%(time, self.symbol)) + func_obj = Function(self.symbol)(time) + pspace_obj = StochasticPSpace(self.symbol, self, self.distribution(time)) + return RandomIndexedSymbol(func_obj, pspace_obj) + +class TransitionMatrixOf(Boolean): + """ + Assumes that the matrix is the transition matrix + of the process. + """ + + def __new__(cls, process, matrix): + if not isinstance(process, DiscreteMarkovChain): + raise ValueError("Currently only DiscreteMarkovChain " + "support TransitionMatrixOf.") + matrix = _matrix_checks(matrix) + return Basic.__new__(cls, process, matrix) + + process = property(lambda self: self.args[0]) + matrix = property(lambda self: self.args[1]) + +class GeneratorMatrixOf(TransitionMatrixOf): + """ + Assumes that the matrix is the generator matrix + of the process. + """ + + def __new__(cls, process, matrix): + if not isinstance(process, ContinuousMarkovChain): + raise ValueError("Currently only ContinuousMarkovChain " + "support GeneratorMatrixOf.") + matrix = _matrix_checks(matrix) + return Basic.__new__(cls, process, matrix) + +class StochasticStateSpaceOf(Boolean): + + def __new__(cls, process, state_space): + if not isinstance(process, (DiscreteMarkovChain, ContinuousMarkovChain)): + raise ValueError("Currently only DiscreteMarkovChain and ContinuousMarkovChain " + "support StochasticStateSpaceOf.") + state_space = _state_converter(state_space) + if isinstance(state_space, Range): + ss_size = ceiling((state_space.stop - state_space.start) / state_space.step) + else: + ss_size = len(state_space) + state_index = Range(ss_size) + return Basic.__new__(cls, process, state_index) + + process = property(lambda self: self.args[0]) + state_index = property(lambda self: self.args[1]) + +class MarkovProcess(StochasticProcess): + """ + Contains methods that handle queries + common to Markov processes. + """ + + @property + def number_of_states(self) -> tUnion[Integer, Symbol]: + """ + The number of states in the Markov Chain. + """ + return _sympify(self.args[2].shape[0]) # type: ignore + + @property + def _state_index(self): + """ + Returns state index as Range. + """ + return self.args[1] + + @classmethod + def _sanity_checks(cls, state_space, trans_probs): + # Try to never have None as state_space or trans_probs. + # This helps a lot if we get it done at the start. + if (state_space is None) and (trans_probs is None): + _n = Dummy('n', integer=True, nonnegative=True) + state_space = _state_converter(Range(_n)) + trans_probs = _matrix_checks(MatrixSymbol('_T', _n, _n)) + + elif state_space is None: + trans_probs = _matrix_checks(trans_probs) + state_space = _state_converter(Range(trans_probs.shape[0])) + + elif trans_probs is None: + state_space = _state_converter(state_space) + if isinstance(state_space, Range): + _n = ceiling((state_space.stop - state_space.start) / state_space.step) + else: + _n = len(state_space) + trans_probs = MatrixSymbol('_T', _n, _n) + + else: + state_space = _state_converter(state_space) + trans_probs = _matrix_checks(trans_probs) + # Range object doesn't want to give a symbolic size + # so we do it ourselves. + if isinstance(state_space, Range): + ss_size = ceiling((state_space.stop - state_space.start) / state_space.step) + else: + ss_size = len(state_space) + if ss_size != trans_probs.shape[0]: + raise ValueError('The size of the state space and the number of ' + 'rows of the transition matrix must be the same.') + + return state_space, trans_probs + + def _extract_information(self, given_condition): + """ + Helper function to extract information, like, + transition matrix/generator matrix, state space, etc. + """ + if isinstance(self, DiscreteMarkovChain): + trans_probs = self.transition_probabilities + state_index = self._state_index + elif isinstance(self, ContinuousMarkovChain): + trans_probs = self.generator_matrix + state_index = self._state_index + if isinstance(given_condition, And): + gcs = given_condition.args + given_condition = S.true + for gc in gcs: + if isinstance(gc, TransitionMatrixOf): + trans_probs = gc.matrix + if isinstance(gc, StochasticStateSpaceOf): + state_index = gc.state_index + if isinstance(gc, Relational): + given_condition = given_condition & gc + if isinstance(given_condition, TransitionMatrixOf): + trans_probs = given_condition.matrix + given_condition = S.true + if isinstance(given_condition, StochasticStateSpaceOf): + state_index = given_condition.state_index + given_condition = S.true + return trans_probs, state_index, given_condition + + def _check_trans_probs(self, trans_probs, row_sum=1): + """ + Helper function for checking the validity of transition + probabilities. + """ + if not isinstance(trans_probs, MatrixSymbol): + rows = trans_probs.tolist() + for row in rows: + if (sum(row) - row_sum) != 0: + raise ValueError("Values in a row must sum to %s. " + "If you are using Float or floats then please use Rational."%(row_sum)) + + def _work_out_state_index(self, state_index, given_condition, trans_probs): + """ + Helper function to extract state space if there + is a random symbol in the given condition. + """ + # if given condition is None, then there is no need to work out + # state_space from random variables + if given_condition != None: + rand_var = list(given_condition.atoms(RandomSymbol) - + given_condition.atoms(RandomIndexedSymbol)) + if len(rand_var) == 1: + state_index = rand_var[0].pspace.set + + # `not None` is `True`. So the old test fails for symbolic sizes. + # Need to build the statement differently. + sym_cond = not self.number_of_states.is_Integer + cond1 = not sym_cond and len(state_index) != trans_probs.shape[0] + if cond1: + raise ValueError("state space is not compatible with the transition probabilities.") + if not isinstance(trans_probs.shape[0], Symbol): + state_index = FiniteSet(*range(trans_probs.shape[0])) + return state_index + + @cacheit + def _preprocess(self, given_condition, evaluate): + """ + Helper function for pre-processing the information. + """ + is_insufficient = False + + if not evaluate: # avoid pre-processing if the result is not to be evaluated + return (True, None, None, None) + + # extracting transition matrix and state space + trans_probs, state_index, given_condition = self._extract_information(given_condition) + + # given_condition does not have sufficient information + # for computations + if trans_probs is None or \ + given_condition is None: + is_insufficient = True + else: + # checking transition probabilities + if isinstance(self, DiscreteMarkovChain): + self._check_trans_probs(trans_probs, row_sum=1) + elif isinstance(self, ContinuousMarkovChain): + self._check_trans_probs(trans_probs, row_sum=0) + + # working out state space + state_index = self._work_out_state_index(state_index, given_condition, trans_probs) + + return is_insufficient, trans_probs, state_index, given_condition + + def replace_with_index(self, condition): + if isinstance(condition, Relational): + lhs, rhs = condition.lhs, condition.rhs + if not isinstance(lhs, RandomIndexedSymbol): + lhs, rhs = rhs, lhs + condition = type(condition)(self.index_of.get(lhs, lhs), + self.index_of.get(rhs, rhs)) + return condition + + def probability(self, condition, given_condition=None, evaluate=True, **kwargs): + """ + Handles probability queries for Markov process. + + Parameters + ========== + + condition: Relational + given_condition: Relational/And + + Returns + ======= + Probability + If the information is not sufficient. + Expr + In all other cases. + + Note + ==== + Any information passed at the time of query overrides + any information passed at the time of object creation like + transition probabilities, state space. + Pass the transition matrix using TransitionMatrixOf, + generator matrix using GeneratorMatrixOf and state space + using StochasticStateSpaceOf in given_condition using & or And. + """ + check, mat, state_index, new_given_condition = \ + self._preprocess(given_condition, evaluate) + + rv = list(condition.atoms(RandomIndexedSymbol)) + symbolic = False + for sym in rv: + if sym.key.is_symbol: + symbolic = True + break + + if check: + return Probability(condition, new_given_condition) + + if isinstance(self, ContinuousMarkovChain): + trans_probs = self.transition_probabilities(mat) + elif isinstance(self, DiscreteMarkovChain): + trans_probs = mat + condition = self.replace_with_index(condition) + given_condition = self.replace_with_index(given_condition) + new_given_condition = self.replace_with_index(new_given_condition) + + if isinstance(condition, Relational): + if isinstance(new_given_condition, And): + gcs = new_given_condition.args + else: + gcs = (new_given_condition, ) + min_key_rv = list(new_given_condition.atoms(RandomIndexedSymbol)) + + if len(min_key_rv): + min_key_rv = min_key_rv[0] + for r in rv: + if min_key_rv.key.is_symbol or r.key.is_symbol: + continue + if min_key_rv.key > r.key: + return Probability(condition) + else: + min_key_rv = None + return Probability(condition) + + if symbolic: + return self._symbolic_probability(condition, new_given_condition, rv, min_key_rv) + + if len(rv) > 1: + rv[0] = condition.lhs + rv[1] = condition.rhs + if rv[0].key < rv[1].key: + rv[0], rv[1] = rv[1], rv[0] + if isinstance(condition, Gt): + condition = Lt(condition.lhs, condition.rhs) + elif isinstance(condition, Lt): + condition = Gt(condition.lhs, condition.rhs) + elif isinstance(condition, Ge): + condition = Le(condition.lhs, condition.rhs) + elif isinstance(condition, Le): + condition = Ge(condition.lhs, condition.rhs) + s = Rational(0, 1) + n = len(self.state_space) + + if isinstance(condition, (Eq, Ne)): + for i in range(0, n): + s += self.probability(Eq(rv[0], i), Eq(rv[1], i)) * self.probability(Eq(rv[1], i), new_given_condition) + return s if isinstance(condition, Eq) else 1 - s + else: + upper = 0 + greater = False + if isinstance(condition, (Ge, Lt)): + upper = 1 + if isinstance(condition, (Ge, Gt)): + greater = True + + for i in range(0, n): + if i <= n//2: + for j in range(0, i + upper): + s += self.probability(Eq(rv[0], i), Eq(rv[1], j)) * self.probability(Eq(rv[1], j), new_given_condition) + else: + s += self.probability(Eq(rv[0], i), new_given_condition) + for j in range(i + upper, n): + s -= self.probability(Eq(rv[0], i), Eq(rv[1], j)) * self.probability(Eq(rv[1], j), new_given_condition) + return s if greater else 1 - s + + rv = rv[0] + states = condition.as_set() + prob, gstate = {}, None + for gc in gcs: + if gc.has(min_key_rv): + if gc.has(Probability): + p, gp = (gc.rhs, gc.lhs) if isinstance(gc.lhs, Probability) \ + else (gc.lhs, gc.rhs) + gr = gp.args[0] + gset = Intersection(gr.as_set(), state_index) + gstate = list(gset)[0] + prob[gset] = p + else: + _, gstate = (gc.lhs.key, gc.rhs) if isinstance(gc.lhs, RandomIndexedSymbol) \ + else (gc.rhs.key, gc.lhs) + + if not all(k in self.index_set for k in (rv.key, min_key_rv.key)): + raise IndexError("The timestamps of the process are not in it's index set.") + states = Intersection(states, state_index) if not isinstance(self.number_of_states, Symbol) else states + for state in Union(states, FiniteSet(gstate)): + if not state.is_Integer or Ge(state, mat.shape[0]) is True: + raise IndexError("No information is available for (%s, %s) in " + "transition probabilities of shape, (%s, %s). " + "State space is zero indexed." + %(gstate, state, mat.shape[0], mat.shape[1])) + if prob: + gstates = Union(*prob.keys()) + if len(gstates) == 1: + gstate = list(gstates)[0] + gprob = list(prob.values())[0] + prob[gstates] = gprob + elif len(gstates) == len(state_index) - 1: + gstate = list(state_index - gstates)[0] + gprob = S.One - sum(prob.values()) + prob[state_index - gstates] = gprob + else: + raise ValueError("Conflicting information.") + else: + gprob = S.One + + if min_key_rv == rv: + return sum(prob[FiniteSet(state)] for state in states) + if isinstance(self, ContinuousMarkovChain): + return gprob * sum(trans_probs(rv.key - min_key_rv.key).__getitem__((gstate, state)) + for state in states) + if isinstance(self, DiscreteMarkovChain): + return gprob * sum((trans_probs**(rv.key - min_key_rv.key)).__getitem__((gstate, state)) + for state in states) + + if isinstance(condition, Not): + expr = condition.args[0] + return S.One - self.probability(expr, given_condition, evaluate, **kwargs) + + if isinstance(condition, And): + compute_later, state2cond, conds = [], {}, condition.args + for expr in conds: + if isinstance(expr, Relational): + ris = list(expr.atoms(RandomIndexedSymbol))[0] + if state2cond.get(ris, None) is None: + state2cond[ris] = S.true + state2cond[ris] &= expr + else: + compute_later.append(expr) + ris = [] + for ri in state2cond: + ris.append(ri) + cset = Intersection(state2cond[ri].as_set(), state_index) + if len(cset) == 0: + return S.Zero + state2cond[ri] = cset.as_relational(ri) + sorted_ris = sorted(ris, key=lambda ri: ri.key) + prod = self.probability(state2cond[sorted_ris[0]], given_condition, evaluate, **kwargs) + for i in range(1, len(sorted_ris)): + ri, prev_ri = sorted_ris[i], sorted_ris[i-1] + if not isinstance(state2cond[ri], Eq): + raise ValueError("The process is in multiple states at %s, unable to determine the probability."%(ri)) + mat_of = TransitionMatrixOf(self, mat) if isinstance(self, DiscreteMarkovChain) else GeneratorMatrixOf(self, mat) + prod *= self.probability(state2cond[ri], state2cond[prev_ri] + & mat_of + & StochasticStateSpaceOf(self, state_index), + evaluate, **kwargs) + for expr in compute_later: + prod *= self.probability(expr, given_condition, evaluate, **kwargs) + return prod + + if isinstance(condition, Or): + return sum(self.probability(expr, given_condition, evaluate, **kwargs) + for expr in condition.args) + + raise NotImplementedError("Mechanism for handling (%s, %s) queries hasn't been " + "implemented yet."%(condition, given_condition)) + + def _symbolic_probability(self, condition, new_given_condition, rv, min_key_rv): + #Function to calculate probability for queries with symbols + if isinstance(condition, Relational): + curr_state = new_given_condition.rhs if isinstance(new_given_condition.lhs, RandomIndexedSymbol) \ + else new_given_condition.lhs + next_state = condition.rhs if isinstance(condition.lhs, RandomIndexedSymbol) \ + else condition.lhs + + if isinstance(condition, (Eq, Ne)): + if isinstance(self, DiscreteMarkovChain): + P = self.transition_probabilities**(rv[0].key - min_key_rv.key) + else: + P = exp(self.generator_matrix*(rv[0].key - min_key_rv.key)) + prob = P[curr_state, next_state] if isinstance(condition, Eq) else 1 - P[curr_state, next_state] + return Piecewise((prob, rv[0].key > min_key_rv.key), (Probability(condition), True)) + else: + upper = 1 + greater = False + if isinstance(condition, (Ge, Lt)): + upper = 0 + if isinstance(condition, (Ge, Gt)): + greater = True + k = Dummy('k') + condition = Eq(condition.lhs, k) if isinstance(condition.lhs, RandomIndexedSymbol)\ + else Eq(condition.rhs, k) + total = Sum(self.probability(condition, new_given_condition), (k, next_state + upper, self.state_space._sup)) + return Piecewise((total, rv[0].key > min_key_rv.key), (Probability(condition), True)) if greater\ + else Piecewise((1 - total, rv[0].key > min_key_rv.key), (Probability(condition), True)) + else: + return Probability(condition, new_given_condition) + + def expectation(self, expr, condition=None, evaluate=True, **kwargs): + """ + Handles expectation queries for markov process. + + Parameters + ========== + + expr: RandomIndexedSymbol, Relational, Logic + Condition for which expectation has to be computed. Must + contain a RandomIndexedSymbol of the process. + condition: Relational, Logic + The given conditions under which computations should be done. + + Returns + ======= + + Expectation + Unevaluated object if computations cannot be done due to + insufficient information. + Expr + In all other cases when the computations are successful. + + Note + ==== + + Any information passed at the time of query overrides + any information passed at the time of object creation like + transition probabilities, state space. + + Pass the transition matrix using TransitionMatrixOf, + generator matrix using GeneratorMatrixOf and state space + using StochasticStateSpaceOf in given_condition using & or And. + """ + + check, mat, state_index, condition = \ + self._preprocess(condition, evaluate) + + if check: + return Expectation(expr, condition) + + rvs = random_symbols(expr) + if isinstance(expr, Expr) and isinstance(condition, Eq) \ + and len(rvs) == 1: + # handle queries similar to E(f(X[i]), Eq(X[i-m], )) + condition=self.replace_with_index(condition) + state_index=self.replace_with_index(state_index) + rv = list(rvs)[0] + lhsg, rhsg = condition.lhs, condition.rhs + if not isinstance(lhsg, RandomIndexedSymbol): + lhsg, rhsg = (rhsg, lhsg) + if rhsg not in state_index: + raise ValueError("%s state is not in the state space."%(rhsg)) + if rv.key < lhsg.key: + raise ValueError("Incorrect given condition is given, expectation " + "time %s < time %s"%(rv.key, rv.key)) + mat_of = TransitionMatrixOf(self, mat) if isinstance(self, DiscreteMarkovChain) else GeneratorMatrixOf(self, mat) + cond = condition & mat_of & \ + StochasticStateSpaceOf(self, state_index) + func = lambda s: self.probability(Eq(rv, s), cond) * expr.subs(rv, self._state_index[s]) + return sum(func(s) for s in state_index) + + raise NotImplementedError("Mechanism for handling (%s, %s) queries hasn't been " + "implemented yet."%(expr, condition)) + +class DiscreteMarkovChain(DiscreteTimeStochasticProcess, MarkovProcess): + """ + Represents a finite discrete time-homogeneous Markov chain. + + This type of Markov Chain can be uniquely characterised by + its (ordered) state space and its one-step transition probability + matrix. + + Parameters + ========== + + sym: + The name given to the Markov Chain + state_space: + Optional, by default, Range(n) + trans_probs: + Optional, by default, MatrixSymbol('_T', n, n) + + Examples + ======== + + >>> from sympy.stats import DiscreteMarkovChain, TransitionMatrixOf, P, E + >>> from sympy import Matrix, MatrixSymbol, Eq, symbols + >>> T = Matrix([[0.5, 0.2, 0.3],[0.2, 0.5, 0.3],[0.2, 0.3, 0.5]]) + >>> Y = DiscreteMarkovChain("Y", [0, 1, 2], T) + >>> YS = DiscreteMarkovChain("Y") + + >>> Y.state_space + {0, 1, 2} + >>> Y.transition_probabilities + Matrix([ + [0.5, 0.2, 0.3], + [0.2, 0.5, 0.3], + [0.2, 0.3, 0.5]]) + >>> TS = MatrixSymbol('T', 3, 3) + >>> P(Eq(YS[3], 2), Eq(YS[1], 1) & TransitionMatrixOf(YS, TS)) + T[0, 2]*T[1, 0] + T[1, 1]*T[1, 2] + T[1, 2]*T[2, 2] + >>> P(Eq(Y[3], 2), Eq(Y[1], 1)).round(2) + 0.36 + + Probabilities will be calculated based on indexes rather + than state names. For example, with the Sunny-Cloudy-Rainy + model with string state names: + + >>> from sympy.core.symbol import Str + >>> Y = DiscreteMarkovChain("Y", [Str('Sunny'), Str('Cloudy'), Str('Rainy')], T) + >>> P(Eq(Y[3], 2), Eq(Y[1], 1)).round(2) + 0.36 + + This gives the same answer as the ``[0, 1, 2]`` state space. + Currently, there is no support for state names within probability + and expectation statements. Here is a work-around using ``Str``: + + >>> P(Eq(Str('Rainy'), Y[3]), Eq(Y[1], Str('Cloudy'))).round(2) + 0.36 + + Symbol state names can also be used: + + >>> sunny, cloudy, rainy = symbols('Sunny, Cloudy, Rainy') + >>> Y = DiscreteMarkovChain("Y", [sunny, cloudy, rainy], T) + >>> P(Eq(Y[3], rainy), Eq(Y[1], cloudy)).round(2) + 0.36 + + Expectations will be calculated as follows: + + >>> E(Y[3], Eq(Y[1], cloudy)) + 0.38*Cloudy + 0.36*Rainy + 0.26*Sunny + + Probability of expressions with multiple RandomIndexedSymbols + can also be calculated provided there is only 1 RandomIndexedSymbol + in the given condition. It is always better to use Rational instead + of floating point numbers for the probabilities in the + transition matrix to avoid errors. + + >>> from sympy import Gt, Le, Rational + >>> T = Matrix([[Rational(5, 10), Rational(3, 10), Rational(2, 10)], [Rational(2, 10), Rational(7, 10), Rational(1, 10)], [Rational(3, 10), Rational(3, 10), Rational(4, 10)]]) + >>> Y = DiscreteMarkovChain("Y", [0, 1, 2], T) + >>> P(Eq(Y[3], Y[1]), Eq(Y[0], 0)).round(3) + 0.409 + >>> P(Gt(Y[3], Y[1]), Eq(Y[0], 0)).round(2) + 0.36 + >>> P(Le(Y[15], Y[10]), Eq(Y[8], 2)).round(7) + 0.6963328 + + Symbolic probability queries are also supported + + >>> a, b, c, d = symbols('a b c d') + >>> T = Matrix([[Rational(1, 10), Rational(4, 10), Rational(5, 10)], [Rational(3, 10), Rational(4, 10), Rational(3, 10)], [Rational(7, 10), Rational(2, 10), Rational(1, 10)]]) + >>> Y = DiscreteMarkovChain("Y", [0, 1, 2], T) + >>> query = P(Eq(Y[a], b), Eq(Y[c], d)) + >>> query.subs({a:10, b:2, c:5, d:1}).round(4) + 0.3096 + >>> P(Eq(Y[10], 2), Eq(Y[5], 1)).evalf().round(4) + 0.3096 + >>> query_gt = P(Gt(Y[a], b), Eq(Y[c], d)) + >>> query_gt.subs({a:21, b:0, c:5, d:0}).evalf().round(5) + 0.64705 + >>> P(Gt(Y[21], 0), Eq(Y[5], 0)).round(5) + 0.64705 + + There is limited support for arbitrarily sized states: + + >>> n = symbols('n', nonnegative=True, integer=True) + >>> T = MatrixSymbol('T', n, n) + >>> Y = DiscreteMarkovChain("Y", trans_probs=T) + >>> Y.state_space + Range(0, n, 1) + >>> query = P(Eq(Y[a], b), Eq(Y[c], d)) + >>> query.subs({a:10, b:2, c:5, d:1}) + (T**5)[1, 2] + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Markov_chain#Discrete-time_Markov_chain + .. [2] https://web.archive.org/web/20201230182007/https://www.dartmouth.edu/~chance/teaching_aids/books_articles/probability_book/Chapter11.pdf + """ + index_set = S.Naturals0 + + def __new__(cls, sym, state_space=None, trans_probs=None): + sym = _symbol_converter(sym) + + state_space, trans_probs = MarkovProcess._sanity_checks(state_space, trans_probs) + + obj = Basic.__new__(cls, sym, state_space, trans_probs) # type: ignore + indices = {} + if isinstance(obj.number_of_states, Integer): + for index, state in enumerate(obj._state_index): + indices[state] = index + obj.index_of = indices + return obj + + @property + def transition_probabilities(self): + """ + Transition probabilities of discrete Markov chain, + either an instance of Matrix or MatrixSymbol. + """ + return self.args[2] + + def communication_classes(self) -> tList[tTuple[tList[Basic], Boolean, Integer]]: + """ + Returns the list of communication classes that partition + the states of the markov chain. + + A communication class is defined to be a set of states + such that every state in that set is reachable from + every other state in that set. Due to its properties + this forms a class in the mathematical sense. + Communication classes are also known as recurrence + classes. + + Returns + ======= + + classes + The ``classes`` are a list of tuples. Each + tuple represents a single communication class + with its properties. The first element in the + tuple is the list of states in the class, the + second element is whether the class is recurrent + and the third element is the period of the + communication class. + + Examples + ======== + + >>> from sympy.stats import DiscreteMarkovChain + >>> from sympy import Matrix + >>> T = Matrix([[0, 1, 0], + ... [1, 0, 0], + ... [1, 0, 0]]) + >>> X = DiscreteMarkovChain('X', [1, 2, 3], T) + >>> classes = X.communication_classes() + >>> for states, is_recurrent, period in classes: + ... states, is_recurrent, period + ([1, 2], True, 2) + ([3], False, 1) + + From this we can see that states ``1`` and ``2`` + communicate, are recurrent and have a period + of 2. We can also see state ``3`` is transient + with a period of 1. + + Notes + ===== + + The algorithm used is of order ``O(n**2)`` where + ``n`` is the number of states in the markov chain. + It uses Tarjan's algorithm to find the classes + themselves and then it uses a breadth-first search + algorithm to find each class's periodicity. + Most of the algorithm's components approach ``O(n)`` + as the matrix becomes more and more sparse. + + References + ========== + + .. [1] https://web.archive.org/web/20220207032113/https://www.columbia.edu/~ww2040/4701Sum07/4701-06-Notes-MCII.pdf + .. [2] https://cecas.clemson.edu/~shierd/Shier/markov.pdf + .. [3] https://www.proquest.com/openview/4adc6a51d8371be5b0e4c7dff287fc70/1?pq-origsite=gscholar&cbl=2026366&diss=y + .. [4] https://www.mathworks.com/help/econ/dtmc.classify.html + """ + n = self.number_of_states + T = self.transition_probabilities + + if isinstance(T, MatrixSymbol): + raise NotImplementedError("Cannot perform the operation with a symbolic matrix.") + + # begin Tarjan's algorithm + V = Range(n) + # don't use state names. Rather use state + # indexes since we use them for matrix + # indexing here and later onward + E = [(i, j) for i in V for j in V if T[i, j] != 0] + classes = strongly_connected_components((V, E)) + # end Tarjan's algorithm + + recurrence = [] + periods = [] + for class_ in classes: + # begin recurrent check (similar to self._check_trans_probs()) + submatrix = T[class_, class_] # get the submatrix with those states + is_recurrent = S.true + rows = submatrix.tolist() + for row in rows: + if (sum(row) - 1) != 0: + is_recurrent = S.false + break + recurrence.append(is_recurrent) + # end recurrent check + + # begin breadth-first search + non_tree_edge_values: tSet[int] = set() + visited = {class_[0]} + newly_visited = {class_[0]} + level = {class_[0]: 0} + current_level = 0 + done = False # imitate a do-while loop + while not done: # runs at most len(class_) times + done = len(visited) == len(class_) + current_level += 1 + + # this loop and the while loop above run a combined len(class_) number of times. + # so this triple nested loop runs through each of the n states once. + for i in newly_visited: + + # the loop below runs len(class_) number of times + # complexity is around about O(n * avg(len(class_))) + newly_visited = {j for j in class_ if T[i, j] != 0} + + new_tree_edges = newly_visited.difference(visited) + for j in new_tree_edges: + level[j] = current_level + + new_non_tree_edges = newly_visited.intersection(visited) + new_non_tree_edge_values = {level[i]-level[j]+1 for j in new_non_tree_edges} + + non_tree_edge_values = non_tree_edge_values.union(new_non_tree_edge_values) + visited = visited.union(new_tree_edges) + + # igcd needs at least 2 arguments + positive_ntev = {val_e for val_e in non_tree_edge_values if val_e > 0} + if len(positive_ntev) == 0: + periods.append(len(class_)) + elif len(positive_ntev) == 1: + periods.append(positive_ntev.pop()) + else: + periods.append(igcd(*positive_ntev)) + # end breadth-first search + + # convert back to the user's state names + classes = [[_sympify(self._state_index[i]) for i in class_] for class_ in classes] + return list(zip(classes, recurrence, map(Integer,periods))) + + def fundamental_matrix(self): + """ + Each entry fundamental matrix can be interpreted as + the expected number of times the chains is in state j + if it started in state i. + + References + ========== + + .. [1] https://lips.cs.princeton.edu/the-fundamental-matrix-of-a-finite-markov-chain/ + + """ + _, _, _, Q = self.decompose() + + if Q.shape[0] > 0: # if non-ergodic + I = eye(Q.shape[0]) + if (I - Q).det() == 0: + raise ValueError("The fundamental matrix doesn't exist.") + return (I - Q).inv().as_immutable() + else: # if ergodic + P = self.transition_probabilities + I = eye(P.shape[0]) + w = self.fixed_row_vector() + W = Matrix([list(w) for i in range(0, P.shape[0])]) + if (I - P + W).det() == 0: + raise ValueError("The fundamental matrix doesn't exist.") + return (I - P + W).inv().as_immutable() + + def absorbing_probabilities(self): + """ + Computes the absorbing probabilities, i.e. + the ij-th entry of the matrix denotes the + probability of Markov chain being absorbed + in state j starting from state i. + """ + _, _, R, _ = self.decompose() + N = self.fundamental_matrix() + if R is None or N is None: + return None + return N*R + + def absorbing_probabilites(self): + sympy_deprecation_warning( + """ + DiscreteMarkovChain.absorbing_probabilites() is deprecated. Use + absorbing_probabilities() instead (note the spelling difference). + """, + deprecated_since_version="1.7", + active_deprecations_target="deprecated-absorbing_probabilites", + ) + return self.absorbing_probabilities() + + def is_regular(self): + tuples = self.communication_classes() + if len(tuples) == 0: + return S.false # not defined for a 0x0 matrix + classes, _, periods = list(zip(*tuples)) + return And(len(classes) == 1, periods[0] == 1) + + def is_ergodic(self): + tuples = self.communication_classes() + if len(tuples) == 0: + return S.false # not defined for a 0x0 matrix + classes, _, _ = list(zip(*tuples)) + return S(len(classes) == 1) + + def is_absorbing_state(self, state): + trans_probs = self.transition_probabilities + if isinstance(trans_probs, ImmutableMatrix) and \ + state < trans_probs.shape[0]: + return S(trans_probs[state, state]) is S.One + + def is_absorbing_chain(self): + states, A, B, C = self.decompose() + r = A.shape[0] + return And(r > 0, A == Identity(r).as_explicit()) + + def stationary_distribution(self, condition_set=False) -> tUnion[ImmutableMatrix, ConditionSet, Lambda]: + r""" + The stationary distribution is any row vector, p, that solves p = pP, + is row stochastic and each element in p must be nonnegative. + That means in matrix form: :math:`(P-I)^T p^T = 0` and + :math:`(1, \dots, 1) p = 1` + where ``P`` is the one-step transition matrix. + + All time-homogeneous Markov Chains with a finite state space + have at least one stationary distribution. In addition, if + a finite time-homogeneous Markov Chain is irreducible, the + stationary distribution is unique. + + Parameters + ========== + + condition_set : bool + If the chain has a symbolic size or transition matrix, + it will return a ``Lambda`` if ``False`` and return a + ``ConditionSet`` if ``True``. + + Examples + ======== + + >>> from sympy.stats import DiscreteMarkovChain + >>> from sympy import Matrix, S + + An irreducible Markov Chain + + >>> T = Matrix([[S(1)/2, S(1)/2, 0], + ... [S(4)/5, S(1)/5, 0], + ... [1, 0, 0]]) + >>> X = DiscreteMarkovChain('X', trans_probs=T) + >>> X.stationary_distribution() + Matrix([[8/13, 5/13, 0]]) + + A reducible Markov Chain + + >>> T = Matrix([[S(1)/2, S(1)/2, 0], + ... [S(4)/5, S(1)/5, 0], + ... [0, 0, 1]]) + >>> X = DiscreteMarkovChain('X', trans_probs=T) + >>> X.stationary_distribution() + Matrix([[8/13 - 8*tau0/13, 5/13 - 5*tau0/13, tau0]]) + + >>> Y = DiscreteMarkovChain('Y') + >>> Y.stationary_distribution() + Lambda((wm, _T), Eq(wm*_T, wm)) + + >>> Y.stationary_distribution(condition_set=True) + ConditionSet(wm, Eq(wm*_T, wm)) + + References + ========== + + .. [1] https://www.probabilitycourse.com/chapter11/11_2_6_stationary_and_limiting_distributions.php + .. [2] https://web.archive.org/web/20210508104430/https://galton.uchicago.edu/~yibi/teaching/stat317/2014/Lectures/Lecture4_6up.pdf + + See Also + ======== + + sympy.stats.DiscreteMarkovChain.limiting_distribution + """ + trans_probs = self.transition_probabilities + n = self.number_of_states + + if n == 0: + return ImmutableMatrix(Matrix([[]])) + + # symbolic matrix version + if isinstance(trans_probs, MatrixSymbol): + wm = MatrixSymbol('wm', 1, n) + if condition_set: + return ConditionSet(wm, Eq(wm * trans_probs, wm)) + else: + return Lambda((wm, trans_probs), Eq(wm * trans_probs, wm)) + + # numeric matrix version + a = Matrix(trans_probs - Identity(n)).T + a[0, 0:n] = ones(1, n) # type: ignore + b = zeros(n, 1) + b[0, 0] = 1 + + soln = list(linsolve((a, b)))[0] + return ImmutableMatrix([soln]) + + def fixed_row_vector(self): + """ + A wrapper for ``stationary_distribution()``. + """ + return self.stationary_distribution() + + @property + def limiting_distribution(self): + """ + The fixed row vector is the limiting + distribution of a discrete Markov chain. + """ + return self.fixed_row_vector() + + def decompose(self) -> tTuple[tList[Basic], ImmutableMatrix, ImmutableMatrix, ImmutableMatrix]: + """ + Decomposes the transition matrix into submatrices with + special properties. + + The transition matrix can be decomposed into 4 submatrices: + - A - the submatrix from recurrent states to recurrent states. + - B - the submatrix from transient to recurrent states. + - C - the submatrix from transient to transient states. + - O - the submatrix of zeros for recurrent to transient states. + + Returns + ======= + + states, A, B, C + ``states`` - a list of state names with the first being + the recurrent states and the last being + the transient states in the order + of the row names of A and then the row names of C. + ``A`` - the submatrix from recurrent states to recurrent states. + ``B`` - the submatrix from transient to recurrent states. + ``C`` - the submatrix from transient to transient states. + + Examples + ======== + + >>> from sympy.stats import DiscreteMarkovChain + >>> from sympy import Matrix, S + + One can decompose this chain for example: + + >>> T = Matrix([[S(1)/2, S(1)/2, 0, 0, 0], + ... [S(2)/5, S(1)/5, S(2)/5, 0, 0], + ... [0, 0, 1, 0, 0], + ... [0, 0, S(1)/2, S(1)/2, 0], + ... [S(1)/2, 0, 0, 0, S(1)/2]]) + >>> X = DiscreteMarkovChain('X', trans_probs=T) + >>> states, A, B, C = X.decompose() + >>> states + [2, 0, 1, 3, 4] + + >>> A # recurrent to recurrent + Matrix([[1]]) + + >>> B # transient to recurrent + Matrix([ + [ 0], + [2/5], + [1/2], + [ 0]]) + + >>> C # transient to transient + Matrix([ + [1/2, 1/2, 0, 0], + [2/5, 1/5, 0, 0], + [ 0, 0, 1/2, 0], + [1/2, 0, 0, 1/2]]) + + This means that state 2 is the only absorbing state + (since A is a 1x1 matrix). B is a 4x1 matrix since + the 4 remaining transient states all merge into reccurent + state 2. And C is the 4x4 matrix that shows how the + transient states 0, 1, 3, 4 all interact. + + See Also + ======== + + sympy.stats.DiscreteMarkovChain.communication_classes + sympy.stats.DiscreteMarkovChain.canonical_form + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Absorbing_Markov_chain + .. [2] https://people.brandeis.edu/~igusa/Math56aS08/Math56a_S08_notes015.pdf + """ + trans_probs = self.transition_probabilities + + classes = self.communication_classes() + r_states = [] + t_states = [] + + for states, recurrent, period in classes: + if recurrent: + r_states += states + else: + t_states += states + + states = r_states + t_states + indexes = [self.index_of[state] for state in states] # type: ignore + + A = Matrix(len(r_states), len(r_states), + lambda i, j: trans_probs[indexes[i], indexes[j]]) + + B = Matrix(len(t_states), len(r_states), + lambda i, j: trans_probs[indexes[len(r_states) + i], indexes[j]]) + + C = Matrix(len(t_states), len(t_states), + lambda i, j: trans_probs[indexes[len(r_states) + i], indexes[len(r_states) + j]]) + + return states, A.as_immutable(), B.as_immutable(), C.as_immutable() + + def canonical_form(self) -> tTuple[tList[Basic], ImmutableMatrix]: + """ + Reorders the one-step transition matrix + so that recurrent states appear first and transient + states appear last. Other representations include inserting + transient states first and recurrent states last. + + Returns + ======= + + states, P_new + ``states`` is the list that describes the order of the + new states in the matrix + so that the ith element in ``states`` is the state of the + ith row of A. + ``P_new`` is the new transition matrix in canonical form. + + Examples + ======== + + >>> from sympy.stats import DiscreteMarkovChain + >>> from sympy import Matrix, S + + You can convert your chain into canonical form: + + >>> T = Matrix([[S(1)/2, S(1)/2, 0, 0, 0], + ... [S(2)/5, S(1)/5, S(2)/5, 0, 0], + ... [0, 0, 1, 0, 0], + ... [0, 0, S(1)/2, S(1)/2, 0], + ... [S(1)/2, 0, 0, 0, S(1)/2]]) + >>> X = DiscreteMarkovChain('X', list(range(1, 6)), trans_probs=T) + >>> states, new_matrix = X.canonical_form() + >>> states + [3, 1, 2, 4, 5] + + >>> new_matrix + Matrix([ + [ 1, 0, 0, 0, 0], + [ 0, 1/2, 1/2, 0, 0], + [2/5, 2/5, 1/5, 0, 0], + [1/2, 0, 0, 1/2, 0], + [ 0, 1/2, 0, 0, 1/2]]) + + The new states are [3, 1, 2, 4, 5] and you can + create a new chain with this and its canonical + form will remain the same (since it is already + in canonical form). + + >>> X = DiscreteMarkovChain('X', states, new_matrix) + >>> states, new_matrix = X.canonical_form() + >>> states + [3, 1, 2, 4, 5] + + >>> new_matrix + Matrix([ + [ 1, 0, 0, 0, 0], + [ 0, 1/2, 1/2, 0, 0], + [2/5, 2/5, 1/5, 0, 0], + [1/2, 0, 0, 1/2, 0], + [ 0, 1/2, 0, 0, 1/2]]) + + This is not limited to absorbing chains: + + >>> T = Matrix([[0, 5, 5, 0, 0], + ... [0, 0, 0, 10, 0], + ... [5, 0, 5, 0, 0], + ... [0, 10, 0, 0, 0], + ... [0, 3, 0, 3, 4]])/10 + >>> X = DiscreteMarkovChain('X', trans_probs=T) + >>> states, new_matrix = X.canonical_form() + >>> states + [1, 3, 0, 2, 4] + + >>> new_matrix + Matrix([ + [ 0, 1, 0, 0, 0], + [ 1, 0, 0, 0, 0], + [ 1/2, 0, 0, 1/2, 0], + [ 0, 0, 1/2, 1/2, 0], + [3/10, 3/10, 0, 0, 2/5]]) + + See Also + ======== + + sympy.stats.DiscreteMarkovChain.communication_classes + sympy.stats.DiscreteMarkovChain.decompose + + References + ========== + + .. [1] https://onlinelibrary.wiley.com/doi/pdf/10.1002/9780470316887.app1 + .. [2] http://www.columbia.edu/~ww2040/6711F12/lect1023big.pdf + """ + states, A, B, C = self.decompose() + O = zeros(A.shape[0], C.shape[1]) + return states, BlockMatrix([[A, O], [B, C]]).as_explicit() + + def sample(self): + """ + Returns + ======= + + sample: iterator object + iterator object containing the sample + + """ + if not isinstance(self.transition_probabilities, (Matrix, ImmutableMatrix)): + raise ValueError("Transition Matrix must be provided for sampling") + Tlist = self.transition_probabilities.tolist() + samps = [random.choice(list(self.state_space))] + yield samps[0] + time = 1 + densities = {} + for state in self.state_space: + states = list(self.state_space) + densities[state] = {states[i]: Tlist[state][i] + for i in range(len(states))} + while time < S.Infinity: + samps.append((next(sample_iter(FiniteRV("_", densities[samps[time - 1]]))))) + yield samps[time] + time += 1 + +class ContinuousMarkovChain(ContinuousTimeStochasticProcess, MarkovProcess): + """ + Represents continuous time Markov chain. + + Parameters + ========== + + sym : Symbol/str + state_space : Set + Optional, by default, S.Reals + gen_mat : Matrix/ImmutableMatrix/MatrixSymbol + Optional, by default, None + + Examples + ======== + + >>> from sympy.stats import ContinuousMarkovChain, P + >>> from sympy import Matrix, S, Eq, Gt + >>> G = Matrix([[-S(1), S(1)], [S(1), -S(1)]]) + >>> C = ContinuousMarkovChain('C', state_space=[0, 1], gen_mat=G) + >>> C.limiting_distribution() + Matrix([[1/2, 1/2]]) + >>> C.state_space + {0, 1} + >>> C.generator_matrix + Matrix([ + [-1, 1], + [ 1, -1]]) + + Probability queries are supported + + >>> P(Eq(C(1.96), 0), Eq(C(0.78), 1)).round(5) + 0.45279 + >>> P(Gt(C(1.7), 0), Eq(C(0.82), 1)).round(5) + 0.58602 + + Probability of expressions with multiple RandomIndexedSymbols + can also be calculated provided there is only 1 RandomIndexedSymbol + in the given condition. It is always better to use Rational instead + of floating point numbers for the probabilities in the + generator matrix to avoid errors. + + >>> from sympy import Gt, Le, Rational + >>> G = Matrix([[-S(1), Rational(1, 10), Rational(9, 10)], [Rational(2, 5), -S(1), Rational(3, 5)], [Rational(1, 2), Rational(1, 2), -S(1)]]) + >>> C = ContinuousMarkovChain('C', state_space=[0, 1, 2], gen_mat=G) + >>> P(Eq(C(3.92), C(1.75)), Eq(C(0.46), 0)).round(5) + 0.37933 + >>> P(Gt(C(3.92), C(1.75)), Eq(C(0.46), 0)).round(5) + 0.34211 + >>> P(Le(C(1.57), C(3.14)), Eq(C(1.22), 1)).round(4) + 0.7143 + + Symbolic probability queries are also supported + + >>> from sympy import symbols + >>> a,b,c,d = symbols('a b c d') + >>> G = Matrix([[-S(1), Rational(1, 10), Rational(9, 10)], [Rational(2, 5), -S(1), Rational(3, 5)], [Rational(1, 2), Rational(1, 2), -S(1)]]) + >>> C = ContinuousMarkovChain('C', state_space=[0, 1, 2], gen_mat=G) + >>> query = P(Eq(C(a), b), Eq(C(c), d)) + >>> query.subs({a:3.65, b:2, c:1.78, d:1}).evalf().round(10) + 0.4002723175 + >>> P(Eq(C(3.65), 2), Eq(C(1.78), 1)).round(10) + 0.4002723175 + >>> query_gt = P(Gt(C(a), b), Eq(C(c), d)) + >>> query_gt.subs({a:43.2, b:0, c:3.29, d:2}).evalf().round(10) + 0.6832579186 + >>> P(Gt(C(43.2), 0), Eq(C(3.29), 2)).round(10) + 0.6832579186 + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Markov_chain#Continuous-time_Markov_chain + .. [2] https://u.math.biu.ac.il/~amirgi/CTMCnotes.pdf + """ + index_set = S.Reals + + def __new__(cls, sym, state_space=None, gen_mat=None): + sym = _symbol_converter(sym) + state_space, gen_mat = MarkovProcess._sanity_checks(state_space, gen_mat) + obj = Basic.__new__(cls, sym, state_space, gen_mat) + indices = {} + if isinstance(obj.number_of_states, Integer): + for index, state in enumerate(obj.state_space): + indices[state] = index + obj.index_of = indices + return obj + + @property + def generator_matrix(self): + return self.args[2] + + @cacheit + def transition_probabilities(self, gen_mat=None): + t = Dummy('t') + if isinstance(gen_mat, (Matrix, ImmutableMatrix)) and \ + gen_mat.is_diagonalizable(): + # for faster computation use diagonalized generator matrix + Q, D = gen_mat.diagonalize() + return Lambda(t, Q*exp(t*D)*Q.inv()) + if gen_mat != None: + return Lambda(t, exp(t*gen_mat)) + + def limiting_distribution(self): + gen_mat = self.generator_matrix + if gen_mat is None: + return None + if isinstance(gen_mat, MatrixSymbol): + wm = MatrixSymbol('wm', 1, gen_mat.shape[0]) + return Lambda((wm, gen_mat), Eq(wm*gen_mat, wm)) + w = IndexedBase('w') + wi = [w[i] for i in range(gen_mat.shape[0])] + wm = Matrix([wi]) + eqs = (wm*gen_mat).tolist()[0] + eqs.append(sum(wi) - 1) + soln = list(linsolve(eqs, wi))[0] + return ImmutableMatrix([soln]) + + +class BernoulliProcess(DiscreteTimeStochasticProcess): + """ + The Bernoulli process consists of repeated + independent Bernoulli process trials with the same parameter `p`. + It's assumed that the probability `p` applies to every + trial and that the outcomes of each trial + are independent of all the rest. Therefore Bernoulli Process + is Discrete State and Discrete Time Stochastic Process. + + Parameters + ========== + + sym : Symbol/str + success : Integer/str + The event which is considered to be success. Default: 1. + failure: Integer/str + The event which is considered to be failure. Default: 0. + p : Real Number between 0 and 1 + Represents the probability of getting success. + + Examples + ======== + + >>> from sympy.stats import BernoulliProcess, P, E + >>> from sympy import Eq, Gt + >>> B = BernoulliProcess("B", p=0.7, success=1, failure=0) + >>> B.state_space + {0, 1} + >>> B.p.round(2) + 0.70 + >>> B.success + 1 + >>> B.failure + 0 + >>> X = B[1] + B[2] + B[3] + >>> P(Eq(X, 0)).round(2) + 0.03 + >>> P(Eq(X, 2)).round(2) + 0.44 + >>> P(Eq(X, 4)).round(2) + 0 + >>> P(Gt(X, 1)).round(2) + 0.78 + >>> P(Eq(B[1], 0) & Eq(B[2], 1) & Eq(B[3], 0) & Eq(B[4], 1)).round(2) + 0.04 + >>> B.joint_distribution(B[1], B[2]) + JointDistributionHandmade(Lambda((B[1], B[2]), Piecewise((0.7, Eq(B[1], 1)), + (0.3, Eq(B[1], 0)), (0, True))*Piecewise((0.7, Eq(B[2], 1)), (0.3, Eq(B[2], 0)), + (0, True)))) + >>> E(2*B[1] + B[2]).round(2) + 2.10 + >>> P(B[1] < 1).round(2) + 0.30 + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Bernoulli_process + .. [2] https://mathcs.clarku.edu/~djoyce/ma217/bernoulli.pdf + + """ + + index_set = S.Naturals0 + + def __new__(cls, sym, p, success=1, failure=0): + _value_check(p >= 0 and p <= 1, 'Value of p must be between 0 and 1.') + sym = _symbol_converter(sym) + p = _sympify(p) + success = _sym_sympify(success) + failure = _sym_sympify(failure) + return Basic.__new__(cls, sym, p, success, failure) + + @property + def symbol(self): + return self.args[0] + + @property + def p(self): + return self.args[1] + + @property + def success(self): + return self.args[2] + + @property + def failure(self): + return self.args[3] + + @property + def state_space(self): + return _set_converter([self.success, self.failure]) + + def distribution(self, key=None): + if key is None: + self._deprecation_warn_distribution() + return BernoulliDistribution(self.p) + return BernoulliDistribution(self.p, self.success, self.failure) + + def simple_rv(self, rv): + return Bernoulli(rv.name, p=self.p, + succ=self.success, fail=self.failure) + + def expectation(self, expr, condition=None, evaluate=True, **kwargs): + """ + Computes expectation. + + Parameters + ========== + + expr : RandomIndexedSymbol, Relational, Logic + Condition for which expectation has to be computed. Must + contain a RandomIndexedSymbol of the process. + condition : Relational, Logic + The given conditions under which computations should be done. + + Returns + ======= + + Expectation of the RandomIndexedSymbol. + + """ + + return _SubstituteRV._expectation(expr, condition, evaluate, **kwargs) + + def probability(self, condition, given_condition=None, evaluate=True, **kwargs): + """ + Computes probability. + + Parameters + ========== + + condition : Relational + Condition for which probability has to be computed. Must + contain a RandomIndexedSymbol of the process. + given_condition : Relational, Logic + The given conditions under which computations should be done. + + Returns + ======= + + Probability of the condition. + + """ + + return _SubstituteRV._probability(condition, given_condition, evaluate, **kwargs) + + def density(self, x): + return Piecewise((self.p, Eq(x, self.success)), + (1 - self.p, Eq(x, self.failure)), + (S.Zero, True)) + +class _SubstituteRV: + """ + Internal class to handle the queries of expectation and probability + by substitution. + """ + + @staticmethod + def _rvindexed_subs(expr, condition=None): + """ + Substitutes the RandomIndexedSymbol with the RandomSymbol with + same name, distribution and probability as RandomIndexedSymbol. + + Parameters + ========== + + expr: RandomIndexedSymbol, Relational, Logic + Condition for which expectation has to be computed. Must + contain a RandomIndexedSymbol of the process. + condition: Relational, Logic + The given conditions under which computations should be done. + + """ + + rvs_expr = random_symbols(expr) + if len(rvs_expr) != 0: + swapdict_expr = {} + for rv in rvs_expr: + if isinstance(rv, RandomIndexedSymbol): + newrv = rv.pspace.process.simple_rv(rv) # substitute with equivalent simple rv + swapdict_expr[rv] = newrv + expr = expr.subs(swapdict_expr) + rvs_cond = random_symbols(condition) + if len(rvs_cond)!=0: + swapdict_cond = {} + for rv in rvs_cond: + if isinstance(rv, RandomIndexedSymbol): + newrv = rv.pspace.process.simple_rv(rv) + swapdict_cond[rv] = newrv + condition = condition.subs(swapdict_cond) + return expr, condition + + @classmethod + def _expectation(self, expr, condition=None, evaluate=True, **kwargs): + """ + Internal method for computing expectation of indexed RV. + + Parameters + ========== + + expr: RandomIndexedSymbol, Relational, Logic + Condition for which expectation has to be computed. Must + contain a RandomIndexedSymbol of the process. + condition: Relational, Logic + The given conditions under which computations should be done. + + Returns + ======= + + Expectation of the RandomIndexedSymbol. + + """ + new_expr, new_condition = self._rvindexed_subs(expr, condition) + + if not is_random(new_expr): + return new_expr + new_pspace = pspace(new_expr) + if new_condition is not None: + new_expr = given(new_expr, new_condition) + if new_expr.is_Add: # As E is Linear + return Add(*[new_pspace.compute_expectation( + expr=arg, evaluate=evaluate, **kwargs) + for arg in new_expr.args]) + return new_pspace.compute_expectation( + new_expr, evaluate=evaluate, **kwargs) + + @classmethod + def _probability(self, condition, given_condition=None, evaluate=True, **kwargs): + """ + Internal method for computing probability of indexed RV + + Parameters + ========== + + condition: Relational + Condition for which probability has to be computed. Must + contain a RandomIndexedSymbol of the process. + given_condition: Relational/And + The given conditions under which computations should be done. + + Returns + ======= + + Probability of the condition. + + """ + new_condition, new_givencondition = self._rvindexed_subs(condition, given_condition) + + if isinstance(new_givencondition, RandomSymbol): + condrv = random_symbols(new_condition) + if len(condrv) == 1 and condrv[0] == new_givencondition: + return BernoulliDistribution(self._probability(new_condition), 0, 1) + + if any(dependent(rv, new_givencondition) for rv in condrv): + return Probability(new_condition, new_givencondition) + else: + return self._probability(new_condition) + + if new_givencondition is not None and \ + not isinstance(new_givencondition, (Relational, Boolean)): + raise ValueError("%s is not a relational or combination of relationals" + % (new_givencondition)) + if new_givencondition == False or new_condition == False: + return S.Zero + if new_condition == True: + return S.One + if not isinstance(new_condition, (Relational, Boolean)): + raise ValueError("%s is not a relational or combination of relationals" + % (new_condition)) + + if new_givencondition is not None: # If there is a condition + # Recompute on new conditional expr + return self._probability(given(new_condition, new_givencondition, **kwargs), **kwargs) + result = pspace(new_condition).probability(new_condition, **kwargs) + if evaluate and hasattr(result, 'doit'): + return result.doit() + else: + return result + +def get_timerv_swaps(expr, condition): + """ + Finds the appropriate interval for each time stamp in expr by parsing + the given condition and returns intervals for each timestamp and + dictionary that maps variable time-stamped Random Indexed Symbol to its + corresponding Random Indexed variable with fixed time stamp. + + Parameters + ========== + + expr: SymPy Expression + Expression containing Random Indexed Symbols with variable time stamps + condition: Relational/Boolean Expression + Expression containing time bounds of variable time stamps in expr + + Examples + ======== + + >>> from sympy.stats.stochastic_process_types import get_timerv_swaps, PoissonProcess + >>> from sympy import symbols, Contains, Interval + >>> x, t, d = symbols('x t d', positive=True) + >>> X = PoissonProcess("X", 3) + >>> get_timerv_swaps(x*X(t), Contains(t, Interval.Lopen(0, 1))) + ([Interval.Lopen(0, 1)], {X(t): X(1)}) + >>> get_timerv_swaps((X(t)**2 + X(d)**2), Contains(t, Interval.Lopen(0, 1)) + ... & Contains(d, Interval.Ropen(1, 4))) # doctest: +SKIP + ([Interval.Ropen(1, 4), Interval.Lopen(0, 1)], {X(d): X(3), X(t): X(1)}) + + Returns + ======= + + intervals: list + List of Intervals/FiniteSet on which each time stamp is defined + rv_swap: dict + Dictionary mapping variable time Random Indexed Symbol to constant time + Random Indexed Variable + + """ + + if not isinstance(condition, (Relational, Boolean)): + raise ValueError("%s is not a relational or combination of relationals" + % (condition)) + expr_syms = list(expr.atoms(RandomIndexedSymbol)) + if isinstance(condition, (And, Or)): + given_cond_args = condition.args + else: # single condition + given_cond_args = (condition, ) + rv_swap = {} + intervals = [] + for expr_sym in expr_syms: + for arg in given_cond_args: + if arg.has(expr_sym.key) and isinstance(expr_sym.key, Symbol): + intv = _set_converter(arg.args[1]) + diff_key = intv._sup - intv._inf + if diff_key == oo: + raise ValueError("%s should have finite bounds" % str(expr_sym.name)) + elif diff_key == S.Zero: # has singleton set + diff_key = intv._sup + rv_swap[expr_sym] = expr_sym.subs({expr_sym.key: diff_key}) + intervals.append(intv) + return intervals, rv_swap + + +class CountingProcess(ContinuousTimeStochasticProcess): + """ + This class handles the common methods of the Counting Processes + such as Poisson, Wiener and Gamma Processes + """ + index_set = _set_converter(Interval(0, oo)) + + @property + def symbol(self): + return self.args[0] + + def expectation(self, expr, condition=None, evaluate=True, **kwargs): + """ + Computes expectation + + Parameters + ========== + + expr: RandomIndexedSymbol, Relational, Logic + Condition for which expectation has to be computed. Must + contain a RandomIndexedSymbol of the process. + condition: Relational, Boolean + The given conditions under which computations should be done, i.e, + the intervals on which each variable time stamp in expr is defined + + Returns + ======= + + Expectation of the given expr + + """ + if condition is not None: + intervals, rv_swap = get_timerv_swaps(expr, condition) + # they are independent when they have non-overlapping intervals + if len(intervals) == 1 or all(Intersection(*intv_comb) == EmptySet + for intv_comb in itertools.combinations(intervals, 2)): + if expr.is_Add: + return Add.fromiter(self.expectation(arg, condition) + for arg in expr.args) + expr = expr.subs(rv_swap) + else: + return Expectation(expr, condition) + + return _SubstituteRV._expectation(expr, evaluate=evaluate, **kwargs) + + def _solve_argwith_tworvs(self, arg): + if arg.args[0].key >= arg.args[1].key or isinstance(arg, Eq): + diff_key = abs(arg.args[0].key - arg.args[1].key) + rv = arg.args[0] + arg = arg.__class__(rv.pspace.process(diff_key), 0) + else: + diff_key = arg.args[1].key - arg.args[0].key + rv = arg.args[1] + arg = arg.__class__(rv.pspace.process(diff_key), 0) + return arg + + def _solve_numerical(self, condition, given_condition=None): + if isinstance(condition, And): + args_list = list(condition.args) + else: + args_list = [condition] + if given_condition is not None: + if isinstance(given_condition, And): + args_list.extend(list(given_condition.args)) + else: + args_list.extend([given_condition]) + # sort the args based on timestamp to get the independent increments in + # each segment using all the condition args as well as given_condition args + args_list = sorted(args_list, key=lambda x: x.args[0].key) + result = [] + cond_args = list(condition.args) if isinstance(condition, And) else [condition] + if args_list[0] in cond_args and not (is_random(args_list[0].args[0]) + and is_random(args_list[0].args[1])): + result.append(_SubstituteRV._probability(args_list[0])) + + if is_random(args_list[0].args[0]) and is_random(args_list[0].args[1]): + arg = self._solve_argwith_tworvs(args_list[0]) + result.append(_SubstituteRV._probability(arg)) + + for i in range(len(args_list) - 1): + curr, nex = args_list[i], args_list[i + 1] + diff_key = nex.args[0].key - curr.args[0].key + working_set = curr.args[0].pspace.process.state_space + if curr.args[1] > nex.args[1]: #impossible condition so return 0 + result.append(0) + break + if isinstance(curr, Eq): + working_set = Intersection(working_set, Interval.Lopen(curr.args[1], oo)) + else: + working_set = Intersection(working_set, curr.as_set()) + if isinstance(nex, Eq): + working_set = Intersection(working_set, Interval(-oo, nex.args[1])) + else: + working_set = Intersection(working_set, nex.as_set()) + if working_set == EmptySet: + rv = Eq(curr.args[0].pspace.process(diff_key), 0) + result.append(_SubstituteRV._probability(rv)) + else: + if working_set.is_finite_set: + if isinstance(curr, Eq) and isinstance(nex, Eq): + rv = Eq(curr.args[0].pspace.process(diff_key), len(working_set)) + result.append(_SubstituteRV._probability(rv)) + elif isinstance(curr, Eq) ^ isinstance(nex, Eq): + result.append(Add.fromiter(_SubstituteRV._probability(Eq( + curr.args[0].pspace.process(diff_key), x)) + for x in range(len(working_set)))) + else: + n = len(working_set) + result.append(Add.fromiter((n - x)*_SubstituteRV._probability(Eq( + curr.args[0].pspace.process(diff_key), x)) for x in range(n))) + else: + result.append(_SubstituteRV._probability( + curr.args[0].pspace.process(diff_key) <= working_set._sup - working_set._inf)) + return Mul.fromiter(result) + + + def probability(self, condition, given_condition=None, evaluate=True, **kwargs): + """ + Computes probability. + + Parameters + ========== + + condition: Relational + Condition for which probability has to be computed. Must + contain a RandomIndexedSymbol of the process. + given_condition: Relational, Boolean + The given conditions under which computations should be done, i.e, + the intervals on which each variable time stamp in expr is defined + + Returns + ======= + + Probability of the condition + + """ + check_numeric = True + if isinstance(condition, (And, Or)): + cond_args = condition.args + else: + cond_args = (condition, ) + # check that condition args are numeric or not + if not all(arg.args[0].key.is_number for arg in cond_args): + check_numeric = False + if given_condition is not None: + check_given_numeric = True + if isinstance(given_condition, (And, Or)): + given_cond_args = given_condition.args + else: + given_cond_args = (given_condition, ) + # check that given condition args are numeric or not + if given_condition.has(Contains): + check_given_numeric = False + # Handle numerical queries + if check_numeric and check_given_numeric: + res = [] + if isinstance(condition, Or): + res.append(Add.fromiter(self._solve_numerical(arg, given_condition) + for arg in condition.args)) + if isinstance(given_condition, Or): + res.append(Add.fromiter(self._solve_numerical(condition, arg) + for arg in given_condition.args)) + if res: + return Add.fromiter(res) + return self._solve_numerical(condition, given_condition) + + # No numeric queries, go by Contains?... then check that all the + # given condition are in form of `Contains` + if not all(arg.has(Contains) for arg in given_cond_args): + raise ValueError("If given condition is passed with `Contains`, then " + "please pass the evaluated condition with its corresponding information " + "in terms of intervals of each time stamp to be passed in given condition.") + + intervals, rv_swap = get_timerv_swaps(condition, given_condition) + # they are independent when they have non-overlapping intervals + if len(intervals) == 1 or all(Intersection(*intv_comb) == EmptySet + for intv_comb in itertools.combinations(intervals, 2)): + if isinstance(condition, And): + return Mul.fromiter(self.probability(arg, given_condition) + for arg in condition.args) + elif isinstance(condition, Or): + return Add.fromiter(self.probability(arg, given_condition) + for arg in condition.args) + condition = condition.subs(rv_swap) + else: + return Probability(condition, given_condition) + if check_numeric: + return self._solve_numerical(condition) + return _SubstituteRV._probability(condition, evaluate=evaluate, **kwargs) + +class PoissonProcess(CountingProcess): + """ + The Poisson process is a counting process. It is usually used in scenarios + where we are counting the occurrences of certain events that appear + to happen at a certain rate, but completely at random. + + Parameters + ========== + + sym : Symbol/str + lamda : Positive number + Rate of the process, ``lambda > 0`` + + Examples + ======== + + >>> from sympy.stats import PoissonProcess, P, E + >>> from sympy import symbols, Eq, Ne, Contains, Interval + >>> X = PoissonProcess("X", lamda=3) + >>> X.state_space + Naturals0 + >>> X.lamda + 3 + >>> t1, t2 = symbols('t1 t2', positive=True) + >>> P(X(t1) < 4) + (9*t1**3/2 + 9*t1**2/2 + 3*t1 + 1)*exp(-3*t1) + >>> P(Eq(X(t1), 2) | Ne(X(t1), 4), Contains(t1, Interval.Ropen(2, 4))) + 1 - 36*exp(-6) + >>> P(Eq(X(t1), 2) & Eq(X(t2), 3), Contains(t1, Interval.Lopen(0, 2)) + ... & Contains(t2, Interval.Lopen(2, 4))) + 648*exp(-12) + >>> E(X(t1)) + 3*t1 + >>> E(X(t1)**2 + 2*X(t2), Contains(t1, Interval.Lopen(0, 1)) + ... & Contains(t2, Interval.Lopen(1, 2))) + 18 + >>> P(X(3) < 1, Eq(X(1), 0)) + exp(-6) + >>> P(Eq(X(4), 3), Eq(X(2), 3)) + exp(-6) + >>> P(X(2) <= 3, X(1) > 1) + 5*exp(-3) + + Merging two Poisson Processes + + >>> Y = PoissonProcess("Y", lamda=4) + >>> Z = X + Y + >>> Z.lamda + 7 + + Splitting a Poisson Process into two independent Poisson Processes + + >>> N, M = Z.split(l1=2, l2=5) + >>> N.lamda, M.lamda + (2, 5) + + References + ========== + + .. [1] https://www.probabilitycourse.com/chapter11/11_0_0_intro.php + .. [2] https://en.wikipedia.org/wiki/Poisson_point_process + + """ + + def __new__(cls, sym, lamda): + _value_check(lamda > 0, 'lamda should be a positive number.') + sym = _symbol_converter(sym) + lamda = _sympify(lamda) + return Basic.__new__(cls, sym, lamda) + + @property + def lamda(self): + return self.args[1] + + @property + def state_space(self): + return S.Naturals0 + + def distribution(self, key): + if isinstance(key, RandomIndexedSymbol): + self._deprecation_warn_distribution() + return PoissonDistribution(self.lamda*key.key) + return PoissonDistribution(self.lamda*key) + + def density(self, x): + return (self.lamda*x.key)**x / factorial(x) * exp(-(self.lamda*x.key)) + + def simple_rv(self, rv): + return Poisson(rv.name, lamda=self.lamda*rv.key) + + def __add__(self, other): + if not isinstance(other, PoissonProcess): + raise ValueError("Only instances of Poisson Process can be merged") + return PoissonProcess(Dummy(self.symbol.name + other.symbol.name), + self.lamda + other.lamda) + + def split(self, l1, l2): + if _sympify(l1 + l2) != self.lamda: + raise ValueError("Sum of l1 and l2 should be %s" % str(self.lamda)) + return PoissonProcess(Dummy("l1"), l1), PoissonProcess(Dummy("l2"), l2) + +class WienerProcess(CountingProcess): + """ + The Wiener process is a real valued continuous-time stochastic process. + In physics it is used to study Brownian motion and it is often also called + Brownian motion due to its historical connection with physical process of the + same name originally observed by Scottish botanist Robert Brown. + + Parameters + ========== + + sym : Symbol/str + + Examples + ======== + + >>> from sympy.stats import WienerProcess, P, E + >>> from sympy import symbols, Contains, Interval + >>> X = WienerProcess("X") + >>> X.state_space + Reals + >>> t1, t2 = symbols('t1 t2', positive=True) + >>> P(X(t1) < 7).simplify() + erf(7*sqrt(2)/(2*sqrt(t1)))/2 + 1/2 + >>> P((X(t1) > 2) | (X(t1) < 4), Contains(t1, Interval.Ropen(2, 4))).simplify() + -erf(1)/2 + erf(2)/2 + 1 + >>> E(X(t1)) + 0 + >>> E(X(t1) + 2*X(t2), Contains(t1, Interval.Lopen(0, 1)) + ... & Contains(t2, Interval.Lopen(1, 2))) + 0 + + References + ========== + + .. [1] https://www.probabilitycourse.com/chapter11/11_4_0_brownian_motion_wiener_process.php + .. [2] https://en.wikipedia.org/wiki/Wiener_process + + """ + def __new__(cls, sym): + sym = _symbol_converter(sym) + return Basic.__new__(cls, sym) + + @property + def state_space(self): + return S.Reals + + def distribution(self, key): + if isinstance(key, RandomIndexedSymbol): + self._deprecation_warn_distribution() + return NormalDistribution(0, sqrt(key.key)) + return NormalDistribution(0, sqrt(key)) + + def density(self, x): + return exp(-x**2/(2*x.key)) / (sqrt(2*pi)*sqrt(x.key)) + + def simple_rv(self, rv): + return Normal(rv.name, 0, sqrt(rv.key)) + + +class GammaProcess(CountingProcess): + r""" + A Gamma process is a random process with independent gamma distributed + increments. It is a pure-jump increasing Levy process. + + Parameters + ========== + + sym : Symbol/str + lamda : Positive number + Jump size of the process, ``lamda > 0`` + gamma : Positive number + Rate of jump arrivals, `\gamma > 0` + + Examples + ======== + + >>> from sympy.stats import GammaProcess, E, P, variance + >>> from sympy import symbols, Contains, Interval, Not + >>> t, d, x, l, g = symbols('t d x l g', positive=True) + >>> X = GammaProcess("X", l, g) + >>> E(X(t)) + g*t/l + >>> variance(X(t)).simplify() + g*t/l**2 + >>> X = GammaProcess('X', 1, 2) + >>> P(X(t) < 1).simplify() + lowergamma(2*t, 1)/gamma(2*t) + >>> P(Not((X(t) < 5) & (X(d) > 3)), Contains(t, Interval.Ropen(2, 4)) & + ... Contains(d, Interval.Lopen(7, 8))).simplify() + -4*exp(-3) + 472*exp(-8)/3 + 1 + >>> E(X(2) + x*E(X(5))) + 10*x + 4 + + References + ========== + + .. [1] https://en.wikipedia.org/wiki/Gamma_process + + """ + def __new__(cls, sym, lamda, gamma): + _value_check(lamda > 0, 'lamda should be a positive number') + _value_check(gamma > 0, 'gamma should be a positive number') + sym = _symbol_converter(sym) + gamma = _sympify(gamma) + lamda = _sympify(lamda) + return Basic.__new__(cls, sym, lamda, gamma) + + @property + def lamda(self): + return self.args[1] + + @property + def gamma(self): + return self.args[2] + + @property + def state_space(self): + return _set_converter(Interval(0, oo)) + + def distribution(self, key): + if isinstance(key, RandomIndexedSymbol): + self._deprecation_warn_distribution() + return GammaDistribution(self.gamma*key.key, 1/self.lamda) + return GammaDistribution(self.gamma*key, 1/self.lamda) + + def density(self, x): + k = self.gamma*x.key + theta = 1/self.lamda + return x**(k - 1) * exp(-x/theta) / (gamma(k)*theta**k) + + def simple_rv(self, rv): + return Gamma(rv.name, self.gamma*rv.key, 1/self.lamda) diff --git a/lib/python3.12/site-packages/sympy/stats/symbolic_multivariate_probability.py b/lib/python3.12/site-packages/sympy/stats/symbolic_multivariate_probability.py new file mode 100644 index 0000000000000000000000000000000000000000..bbe8776e58e82489e29734cea48c9138bc512f34 --- /dev/null +++ b/lib/python3.12/site-packages/sympy/stats/symbolic_multivariate_probability.py @@ -0,0 +1,308 @@ +import itertools + +from sympy.core.add import Add +from sympy.core.expr import Expr +from sympy.core.function import expand as _expand +from sympy.core.mul import Mul +from sympy.core.singleton import S +from sympy.matrices.exceptions import ShapeError +from sympy.matrices.expressions.matexpr import MatrixExpr +from sympy.matrices.expressions.matmul import MatMul +from sympy.matrices.expressions.special import ZeroMatrix +from sympy.stats.rv import RandomSymbol, is_random +from sympy.core.sympify import _sympify +from sympy.stats.symbolic_probability import Variance, Covariance, Expectation + + +class ExpectationMatrix(Expectation, MatrixExpr): + """ + Expectation of a random matrix expression. + + Examples + ======== + + >>> from sympy.stats import ExpectationMatrix, Normal + >>> from sympy.stats.rv import RandomMatrixSymbol + >>> from sympy import symbols, MatrixSymbol, Matrix + >>> k = symbols("k") + >>> A, B = MatrixSymbol("A", k, k), MatrixSymbol("B", k, k) + >>> X, Y = RandomMatrixSymbol("X", k, 1), RandomMatrixSymbol("Y", k, 1) + >>> ExpectationMatrix(X) + ExpectationMatrix(X) + >>> ExpectationMatrix(A*X).shape + (k, 1) + + To expand the expectation in its expression, use ``expand()``: + + >>> ExpectationMatrix(A*X + B*Y).expand() + A*ExpectationMatrix(X) + B*ExpectationMatrix(Y) + >>> ExpectationMatrix((X + Y)*(X - Y).T).expand() + ExpectationMatrix(X*X.T) - ExpectationMatrix(X*Y.T) + ExpectationMatrix(Y*X.T) - ExpectationMatrix(Y*Y.T) + + To evaluate the ``ExpectationMatrix``, use ``doit()``: + + >>> N11, N12 = Normal('N11', 11, 1), Normal('N12', 12, 1) + >>> N21, N22 = Normal('N21', 21, 1), Normal('N22', 22, 1) + >>> M11, M12 = Normal('M11', 1, 1), Normal('M12', 2, 1) + >>> M21, M22 = Normal('M21', 3, 1), Normal('M22', 4, 1) + >>> x1 = Matrix([[N11, N12], [N21, N22]]) + >>> x2 = Matrix([[M11, M12], [M21, M22]]) + >>> ExpectationMatrix(x1 + x2).doit() + Matrix([ + [12, 14], + [24, 26]]) + + """ + def __new__(cls, expr, condition=None): + expr = _sympify(expr) + if condition is None: + if not is_random(expr): + return expr + obj = Expr.__new__(cls, expr) + else: + condition = _sympify(condition) + obj = Expr.__new__(cls, expr, condition) + + obj._shape = expr.shape + obj._condition = condition + return obj + + @property + def shape(self): + return self._shape + + def expand(self, **hints): + expr = self.args[0] + condition = self._condition + if not is_random(expr): + return expr + + if isinstance(expr, Add): + return Add.fromiter(Expectation(a, condition=condition).expand() + for a in expr.args) + + expand_expr = _expand(expr) + if isinstance(expand_expr, Add): + return Add.fromiter(Expectation(a, condition=condition).expand() + for a in expand_expr.args) + + elif isinstance(expr, (Mul, MatMul)): + rv = [] + nonrv = [] + postnon = [] + + for a in expr.args: + if is_random(a): + if rv: + rv.extend(postnon) + else: + nonrv.extend(postnon) + postnon = [] + rv.append(a) + elif a.is_Matrix: + postnon.append(a) + else: + nonrv.append(a) + + # In order to avoid infinite-looping (MatMul may call .doit() again), + # do not rebuild + if len(nonrv) == 0: + return self + return Mul.fromiter(nonrv)*Expectation(Mul.fromiter(rv), + condition=condition)*Mul.fromiter(postnon) + + return self + +class VarianceMatrix(Variance, MatrixExpr): + """ + Variance of a random matrix probability expression. Also known as + Covariance matrix, auto-covariance matrix, dispersion matrix, + or variance-covariance matrix. + + Examples + ======== + + >>> from sympy.stats import VarianceMatrix + >>> from sympy.stats.rv import RandomMatrixSymbol + >>> from sympy import symbols, MatrixSymbol + >>> k = symbols("k") + >>> A, B = MatrixSymbol("A", k, k), MatrixSymbol("B", k, k) + >>> X, Y = RandomMatrixSymbol("X", k, 1), RandomMatrixSymbol("Y", k, 1) + >>> VarianceMatrix(X) + VarianceMatrix(X) + >>> VarianceMatrix(X).shape + (k, k) + + To expand the variance in its expression, use ``expand()``: + + >>> VarianceMatrix(A*X).expand() + A*VarianceMatrix(X)*A.T + >>> VarianceMatrix(A*X + B*Y).expand() + 2*A*CrossCovarianceMatrix(X, Y)*B.T + A*VarianceMatrix(X)*A.T + B*VarianceMatrix(Y)*B.T + """ + def __new__(cls, arg, condition=None): + arg = _sympify(arg) + + if 1 not in arg.shape: + raise ShapeError("Expression is not a vector") + + shape = (arg.shape[0], arg.shape[0]) if arg.shape[1] == 1 else (arg.shape[1], arg.shape[1]) + + if condition: + obj = Expr.__new__(cls, arg, condition) + else: + obj = Expr.__new__(cls, arg) + + obj._shape = shape + obj._condition = condition + return obj + + @property + def shape(self): + return self._shape + + def expand(self, **hints): + arg = self.args[0] + condition = self._condition + + if not is_random(arg): + return ZeroMatrix(*self.shape) + + if isinstance(arg, RandomSymbol): + return self + elif isinstance(arg, Add): + rv = [] + for a in arg.args: + if is_random(a): + rv.append(a) + variances = Add(*(Variance(xv, condition).expand() for xv in rv)) + map_to_covar = lambda x: 2*Covariance(*x, condition=condition).expand() + covariances = Add(*map(map_to_covar, itertools.combinations(rv, 2))) + return variances + covariances + elif isinstance(arg, (Mul, MatMul)): + nonrv = [] + rv = [] + for a in arg.args: + if is_random(a): + rv.append(a) + else: + nonrv.append(a) + if len(rv) == 0: + return ZeroMatrix(*self.shape) + # Avoid possible infinite loops with MatMul: + if len(nonrv) == 0: + return self + # Variance of many multiple matrix products is not implemented: + if len(rv) > 1: + return self + return Mul.fromiter(nonrv)*Variance(Mul.fromiter(rv), + condition)*(Mul.fromiter(nonrv)).transpose() + + # this expression contains a RandomSymbol somehow: + return self + +class CrossCovarianceMatrix(Covariance, MatrixExpr): + """ + Covariance of a random matrix probability expression. + + Examples + ======== + + >>> from sympy.stats import CrossCovarianceMatrix + >>> from sympy.stats.rv import RandomMatrixSymbol + >>> from sympy import symbols, MatrixSymbol + >>> k = symbols("k") + >>> A, B = MatrixSymbol("A", k, k), MatrixSymbol("B", k, k) + >>> C, D = MatrixSymbol("C", k, k), MatrixSymbol("D", k, k) + >>> X, Y = RandomMatrixSymbol("X", k, 1), RandomMatrixSymbol("Y", k, 1) + >>> Z, W = RandomMatrixSymbol("Z", k, 1), RandomMatrixSymbol("W", k, 1) + >>> CrossCovarianceMatrix(X, Y) + CrossCovarianceMatrix(X, Y) + >>> CrossCovarianceMatrix(X, Y).shape + (k, k) + + To expand the covariance in its expression, use ``expand()``: + + >>> CrossCovarianceMatrix(X + Y, Z).expand() + CrossCovarianceMatrix(X, Z) + CrossCovarianceMatrix(Y, Z) + >>> CrossCovarianceMatrix(A*X, Y).expand() + A*CrossCovarianceMatrix(X, Y) + >>> CrossCovarianceMatrix(A*X, B.T*Y).expand() + A*CrossCovarianceMatrix(X, Y)*B + >>> CrossCovarianceMatrix(A*X + B*Y, C.T*Z + D.T*W).expand() + A*CrossCovarianceMatrix(X, W)*D + A*CrossCovarianceMatrix(X, Z)*C + B*CrossCovarianceMatrix(Y, W)*D + B*CrossCovarianceMatrix(Y, Z)*C + + """ + def __new__(cls, arg1, arg2, condition=None): + arg1 = _sympify(arg1) + arg2 = _sympify(arg2) + + if (1 not in arg1.shape) or (1 not in arg2.shape) or (arg1.shape[1] != arg2.shape[1]): + raise ShapeError("Expression is not a vector") + + shape = (arg1.shape[0], arg2.shape[0]) if arg1.shape[1] == 1 and arg2.shape[1] == 1 \ + else (1, 1) + + if condition: + obj = Expr.__new__(cls, arg1, arg2, condition) + else: + obj = Expr.__new__(cls, arg1, arg2) + + obj._shape = shape + obj._condition = condition + return obj + + @property + def shape(self): + return self._shape + + def expand(self, **hints): + arg1 = self.args[0] + arg2 = self.args[1] + condition = self._condition + + if arg1 == arg2: + return VarianceMatrix(arg1, condition).expand() + + if not is_random(arg1) or not is_random(arg2): + return ZeroMatrix(*self.shape) + + if isinstance(arg1, RandomSymbol) and isinstance(arg2, RandomSymbol): + return CrossCovarianceMatrix(arg1, arg2, condition) + + coeff_rv_list1 = self._expand_single_argument(arg1.expand()) + coeff_rv_list2 = self._expand_single_argument(arg2.expand()) + + addends = [a*CrossCovarianceMatrix(r1, r2, condition=condition)*b.transpose() + for (a, r1) in coeff_rv_list1 for (b, r2) in coeff_rv_list2] + return Add.fromiter(addends) + + @classmethod + def _expand_single_argument(cls, expr): + # return (coefficient, random_symbol) pairs: + if isinstance(expr, RandomSymbol): + return [(S.One, expr)] + elif isinstance(expr, Add): + outval = [] + for a in expr.args: + if isinstance(a, (Mul, MatMul)): + outval.append(cls._get_mul_nonrv_rv_tuple(a)) + elif is_random(a): + outval.append((S.One, a)) + + return outval + elif isinstance(expr, (Mul, MatMul)): + return [cls._get_mul_nonrv_rv_tuple(expr)] + elif is_random(expr): + return [(S.One, expr)] + + @classmethod + def _get_mul_nonrv_rv_tuple(cls, m): + rv = [] + nonrv = [] + for a in m.args: + if is_random(a): + rv.append(a) + else: + nonrv.append(a) + return (Mul.fromiter(nonrv), Mul.fromiter(rv)) diff --git a/lib/python3.12/site-packages/sympy/stats/symbolic_probability.py b/lib/python3.12/site-packages/sympy/stats/symbolic_probability.py new file mode 100644 index 0000000000000000000000000000000000000000..5d0b971a8691f82de15258d4c460129059eaf436 --- /dev/null +++ b/lib/python3.12/site-packages/sympy/stats/symbolic_probability.py @@ -0,0 +1,698 @@ +import itertools +from sympy.concrete.summations import Sum +from sympy.core.add import Add +from sympy.core.expr import Expr +from sympy.core.function import expand as _expand +from sympy.core.mul import Mul +from sympy.core.relational import Eq +from sympy.core.singleton import S +from sympy.core.symbol import Symbol +from sympy.integrals.integrals import Integral +from sympy.logic.boolalg import Not +from sympy.core.parameters import global_parameters +from sympy.core.sorting import default_sort_key +from sympy.core.sympify import _sympify +from sympy.core.relational import Relational +from sympy.logic.boolalg import Boolean +from sympy.stats import variance, covariance +from sympy.stats.rv import (RandomSymbol, pspace, dependent, + given, sampling_E, RandomIndexedSymbol, is_random, + PSpace, sampling_P, random_symbols) + +__all__ = ['Probability', 'Expectation', 'Variance', 'Covariance'] + + +@is_random.register(Expr) +def _(x): + atoms = x.free_symbols + if len(atoms) == 1 and next(iter(atoms)) == x: + return False + return any(is_random(i) for i in atoms) + +@is_random.register(RandomSymbol) # type: ignore +def _(x): + return True + + +class Probability(Expr): + """ + Symbolic expression for the probability. + + Examples + ======== + + >>> from sympy.stats import Probability, Normal + >>> from sympy import Integral + >>> X = Normal("X", 0, 1) + >>> prob = Probability(X > 1) + >>> prob + Probability(X > 1) + + Integral representation: + + >>> prob.rewrite(Integral) + Integral(sqrt(2)*exp(-_z**2/2)/(2*sqrt(pi)), (_z, 1, oo)) + + Evaluation of the integral: + + >>> prob.evaluate_integral() + sqrt(2)*(-sqrt(2)*sqrt(pi)*erf(sqrt(2)/2) + sqrt(2)*sqrt(pi))/(4*sqrt(pi)) + """ + + is_commutative = True + + def __new__(cls, prob, condition=None, **kwargs): + prob = _sympify(prob) + if condition is None: + obj = Expr.__new__(cls, prob) + else: + condition = _sympify(condition) + obj = Expr.__new__(cls, prob, condition) + obj._condition = condition + return obj + + def doit(self, **hints): + condition = self.args[0] + given_condition = self._condition + numsamples = hints.get('numsamples', False) + evaluate = hints.get('evaluate', True) + + if isinstance(condition, Not): + return S.One - self.func(condition.args[0], given_condition, + evaluate=evaluate).doit(**hints) + + if condition.has(RandomIndexedSymbol): + return pspace(condition).probability(condition, given_condition, + evaluate=evaluate) + + if isinstance(given_condition, RandomSymbol): + condrv = random_symbols(condition) + if len(condrv) == 1 and condrv[0] == given_condition: + from sympy.stats.frv_types import BernoulliDistribution + return BernoulliDistribution(self.func(condition).doit(**hints), 0, 1) + if any(dependent(rv, given_condition) for rv in condrv): + return Probability(condition, given_condition) + else: + return Probability(condition).doit() + + if given_condition is not None and \ + not isinstance(given_condition, (Relational, Boolean)): + raise ValueError("%s is not a relational or combination of relationals" + % (given_condition)) + + if given_condition == False or condition is S.false: + return S.Zero + if not isinstance(condition, (Relational, Boolean)): + raise ValueError("%s is not a relational or combination of relationals" + % (condition)) + if condition is S.true: + return S.One + + if numsamples: + return sampling_P(condition, given_condition, numsamples=numsamples) + if given_condition is not None: # If there is a condition + # Recompute on new conditional expr + return Probability(given(condition, given_condition)).doit() + + # Otherwise pass work off to the ProbabilitySpace + if pspace(condition) == PSpace(): + return Probability(condition, given_condition) + + result = pspace(condition).probability(condition) + if hasattr(result, 'doit') and evaluate: + return result.doit() + else: + return result + + def _eval_rewrite_as_Integral(self, arg, condition=None, **kwargs): + return self.func(arg, condition=condition).doit(evaluate=False) + + _eval_rewrite_as_Sum = _eval_rewrite_as_Integral + + def evaluate_integral(self): + return self.rewrite(Integral).doit() + + +class Expectation(Expr): + """ + Symbolic expression for the expectation. + + Examples + ======== + + >>> from sympy.stats import Expectation, Normal, Probability, Poisson + >>> from sympy import symbols, Integral, Sum + >>> mu = symbols("mu") + >>> sigma = symbols("sigma", positive=True) + >>> X = Normal("X", mu, sigma) + >>> Expectation(X) + Expectation(X) + >>> Expectation(X).evaluate_integral().simplify() + mu + + To get the integral expression of the expectation: + + >>> Expectation(X).rewrite(Integral) + Integral(sqrt(2)*X*exp(-(X - mu)**2/(2*sigma**2))/(2*sqrt(pi)*sigma), (X, -oo, oo)) + + The same integral expression, in more abstract terms: + + >>> Expectation(X).rewrite(Probability) + Integral(x*Probability(Eq(X, x)), (x, -oo, oo)) + + To get the Summation expression of the expectation for discrete random variables: + + >>> lamda = symbols('lamda', positive=True) + >>> Z = Poisson('Z', lamda) + >>> Expectation(Z).rewrite(Sum) + Sum(Z*lamda**Z*exp(-lamda)/factorial(Z), (Z, 0, oo)) + + This class is aware of some properties of the expectation: + + >>> from sympy.abc import a + >>> Expectation(a*X) + Expectation(a*X) + >>> Y = Normal("Y", 1, 2) + >>> Expectation(X + Y) + Expectation(X + Y) + + To expand the ``Expectation`` into its expression, use ``expand()``: + + >>> Expectation(X + Y).expand() + Expectation(X) + Expectation(Y) + >>> Expectation(a*X + Y).expand() + a*Expectation(X) + Expectation(Y) + >>> Expectation(a*X + Y) + Expectation(a*X + Y) + >>> Expectation((X + Y)*(X - Y)).expand() + Expectation(X**2) - Expectation(Y**2) + + To evaluate the ``Expectation``, use ``doit()``: + + >>> Expectation(X + Y).doit() + mu + 1 + >>> Expectation(X + Expectation(Y + Expectation(2*X))).doit() + 3*mu + 1 + + To prevent evaluating nested ``Expectation``, use ``doit(deep=False)`` + + >>> Expectation(X + Expectation(Y)).doit(deep=False) + mu + Expectation(Expectation(Y)) + >>> Expectation(X + Expectation(Y + Expectation(2*X))).doit(deep=False) + mu + Expectation(Expectation(Expectation(2*X) + Y)) + + """ + + def __new__(cls, expr, condition=None, **kwargs): + expr = _sympify(expr) + if expr.is_Matrix: + from sympy.stats.symbolic_multivariate_probability import ExpectationMatrix + return ExpectationMatrix(expr, condition) + if condition is None: + if not is_random(expr): + return expr + obj = Expr.__new__(cls, expr) + else: + condition = _sympify(condition) + obj = Expr.__new__(cls, expr, condition) + obj._condition = condition + return obj + + def _eval_is_commutative(self): + return(self.args[0].is_commutative) + + def expand(self, **hints): + expr = self.args[0] + condition = self._condition + + if not is_random(expr): + return expr + + if isinstance(expr, Add): + return Add.fromiter(Expectation(a, condition=condition).expand() + for a in expr.args) + + expand_expr = _expand(expr) + if isinstance(expand_expr, Add): + return Add.fromiter(Expectation(a, condition=condition).expand() + for a in expand_expr.args) + + elif isinstance(expr, Mul): + rv = [] + nonrv = [] + for a in expr.args: + if is_random(a): + rv.append(a) + else: + nonrv.append(a) + return Mul.fromiter(nonrv)*Expectation(Mul.fromiter(rv), condition=condition) + + return self + + def doit(self, **hints): + deep = hints.get('deep', True) + condition = self._condition + expr = self.args[0] + numsamples = hints.get('numsamples', False) + evaluate = hints.get('evaluate', True) + + if deep: + expr = expr.doit(**hints) + + if not is_random(expr) or isinstance(expr, Expectation): # expr isn't random? + return expr + if numsamples: # Computing by monte carlo sampling? + evalf = hints.get('evalf', True) + return sampling_E(expr, condition, numsamples=numsamples, evalf=evalf) + + if expr.has(RandomIndexedSymbol): + return pspace(expr).compute_expectation(expr, condition) + + # Create new expr and recompute E + if condition is not None: # If there is a condition + return self.func(given(expr, condition)).doit(**hints) + + # A few known statements for efficiency + + if expr.is_Add: # We know that E is Linear + return Add(*[self.func(arg, condition).doit(**hints) + if not isinstance(arg, Expectation) else self.func(arg, condition) + for arg in expr.args]) + if expr.is_Mul: + if expr.atoms(Expectation): + return expr + + if pspace(expr) == PSpace(): + return self.func(expr) + # Otherwise case is simple, pass work off to the ProbabilitySpace + result = pspace(expr).compute_expectation(expr, evaluate=evaluate) + if hasattr(result, 'doit') and evaluate: + return result.doit(**hints) + else: + return result + + + def _eval_rewrite_as_Probability(self, arg, condition=None, **kwargs): + rvs = arg.atoms(RandomSymbol) + if len(rvs) > 1: + raise NotImplementedError() + if len(rvs) == 0: + return arg + + rv = rvs.pop() + if rv.pspace is None: + raise ValueError("Probability space not known") + + symbol = rv.symbol + if symbol.name[0].isupper(): + symbol = Symbol(symbol.name.lower()) + else : + symbol = Symbol(symbol.name + "_1") + + if rv.pspace.is_Continuous: + return Integral(arg.replace(rv, symbol)*Probability(Eq(rv, symbol), condition), (symbol, rv.pspace.domain.set.inf, rv.pspace.domain.set.sup)) + else: + if rv.pspace.is_Finite: + raise NotImplementedError + else: + return Sum(arg.replace(rv, symbol)*Probability(Eq(rv, symbol), condition), (symbol, rv.pspace.domain.set.inf, rv.pspace.set.sup)) + + def _eval_rewrite_as_Integral(self, arg, condition=None, evaluate=False, **kwargs): + return self.func(arg, condition=condition).doit(deep=False, evaluate=evaluate) + + _eval_rewrite_as_Sum = _eval_rewrite_as_Integral # For discrete this will be Sum + + def evaluate_integral(self): + return self.rewrite(Integral).doit() + + evaluate_sum = evaluate_integral + +class Variance(Expr): + """ + Symbolic expression for the variance. + + Examples + ======== + + >>> from sympy import symbols, Integral + >>> from sympy.stats import Normal, Expectation, Variance, Probability + >>> mu = symbols("mu", positive=True) + >>> sigma = symbols("sigma", positive=True) + >>> X = Normal("X", mu, sigma) + >>> Variance(X) + Variance(X) + >>> Variance(X).evaluate_integral() + sigma**2 + + Integral representation of the underlying calculations: + + >>> Variance(X).rewrite(Integral) + Integral(sqrt(2)*(X - Integral(sqrt(2)*X*exp(-(X - mu)**2/(2*sigma**2))/(2*sqrt(pi)*sigma), (X, -oo, oo)))**2*exp(-(X - mu)**2/(2*sigma**2))/(2*sqrt(pi)*sigma), (X, -oo, oo)) + + Integral representation, without expanding the PDF: + + >>> Variance(X).rewrite(Probability) + -Integral(x*Probability(Eq(X, x)), (x, -oo, oo))**2 + Integral(x**2*Probability(Eq(X, x)), (x, -oo, oo)) + + Rewrite the variance in terms of the expectation + + >>> Variance(X).rewrite(Expectation) + -Expectation(X)**2 + Expectation(X**2) + + Some transformations based on the properties of the variance may happen: + + >>> from sympy.abc import a + >>> Y = Normal("Y", 0, 1) + >>> Variance(a*X) + Variance(a*X) + + To expand the variance in its expression, use ``expand()``: + + >>> Variance(a*X).expand() + a**2*Variance(X) + >>> Variance(X + Y) + Variance(X + Y) + >>> Variance(X + Y).expand() + 2*Covariance(X, Y) + Variance(X) + Variance(Y) + + """ + def __new__(cls, arg, condition=None, **kwargs): + arg = _sympify(arg) + + if arg.is_Matrix: + from sympy.stats.symbolic_multivariate_probability import VarianceMatrix + return VarianceMatrix(arg, condition) + if condition is None: + obj = Expr.__new__(cls, arg) + else: + condition = _sympify(condition) + obj = Expr.__new__(cls, arg, condition) + obj._condition = condition + return obj + + def _eval_is_commutative(self): + return self.args[0].is_commutative + + def expand(self, **hints): + arg = self.args[0] + condition = self._condition + + if not is_random(arg): + return S.Zero + + if isinstance(arg, RandomSymbol): + return self + elif isinstance(arg, Add): + rv = [] + for a in arg.args: + if is_random(a): + rv.append(a) + variances = Add(*(Variance(xv, condition).expand() for xv in rv)) + map_to_covar = lambda x: 2*Covariance(*x, condition=condition).expand() + covariances = Add(*map(map_to_covar, itertools.combinations(rv, 2))) + return variances + covariances + elif isinstance(arg, Mul): + nonrv = [] + rv = [] + for a in arg.args: + if is_random(a): + rv.append(a) + else: + nonrv.append(a**2) + if len(rv) == 0: + return S.Zero + return Mul.fromiter(nonrv)*Variance(Mul.fromiter(rv), condition) + + # this expression contains a RandomSymbol somehow: + return self + + def _eval_rewrite_as_Expectation(self, arg, condition=None, **kwargs): + e1 = Expectation(arg**2, condition) + e2 = Expectation(arg, condition)**2 + return e1 - e2 + + def _eval_rewrite_as_Probability(self, arg, condition=None, **kwargs): + return self.rewrite(Expectation).rewrite(Probability) + + def _eval_rewrite_as_Integral(self, arg, condition=None, **kwargs): + return variance(self.args[0], self._condition, evaluate=False) + + _eval_rewrite_as_Sum = _eval_rewrite_as_Integral + + def evaluate_integral(self): + return self.rewrite(Integral).doit() + + +class Covariance(Expr): + """ + Symbolic expression for the covariance. + + Examples + ======== + + >>> from sympy.stats import Covariance + >>> from sympy.stats import Normal + >>> X = Normal("X", 3, 2) + >>> Y = Normal("Y", 0, 1) + >>> Z = Normal("Z", 0, 1) + >>> W = Normal("W", 0, 1) + >>> cexpr = Covariance(X, Y) + >>> cexpr + Covariance(X, Y) + + Evaluate the covariance, `X` and `Y` are independent, + therefore zero is the result: + + >>> cexpr.evaluate_integral() + 0 + + Rewrite the covariance expression in terms of expectations: + + >>> from sympy.stats import Expectation + >>> cexpr.rewrite(Expectation) + Expectation(X*Y) - Expectation(X)*Expectation(Y) + + In order to expand the argument, use ``expand()``: + + >>> from sympy.abc import a, b, c, d + >>> Covariance(a*X + b*Y, c*Z + d*W) + Covariance(a*X + b*Y, c*Z + d*W) + >>> Covariance(a*X + b*Y, c*Z + d*W).expand() + a*c*Covariance(X, Z) + a*d*Covariance(W, X) + b*c*Covariance(Y, Z) + b*d*Covariance(W, Y) + + This class is aware of some properties of the covariance: + + >>> Covariance(X, X).expand() + Variance(X) + >>> Covariance(a*X, b*Y).expand() + a*b*Covariance(X, Y) + """ + + def __new__(cls, arg1, arg2, condition=None, **kwargs): + arg1 = _sympify(arg1) + arg2 = _sympify(arg2) + + if arg1.is_Matrix or arg2.is_Matrix: + from sympy.stats.symbolic_multivariate_probability import CrossCovarianceMatrix + return CrossCovarianceMatrix(arg1, arg2, condition) + + if kwargs.pop('evaluate', global_parameters.evaluate): + arg1, arg2 = sorted([arg1, arg2], key=default_sort_key) + + if condition is None: + obj = Expr.__new__(cls, arg1, arg2) + else: + condition = _sympify(condition) + obj = Expr.__new__(cls, arg1, arg2, condition) + obj._condition = condition + return obj + + def _eval_is_commutative(self): + return self.args[0].is_commutative + + def expand(self, **hints): + arg1 = self.args[0] + arg2 = self.args[1] + condition = self._condition + + if arg1 == arg2: + return Variance(arg1, condition).expand() + + if not is_random(arg1): + return S.Zero + if not is_random(arg2): + return S.Zero + + arg1, arg2 = sorted([arg1, arg2], key=default_sort_key) + + if isinstance(arg1, RandomSymbol) and isinstance(arg2, RandomSymbol): + return Covariance(arg1, arg2, condition) + + coeff_rv_list1 = self._expand_single_argument(arg1.expand()) + coeff_rv_list2 = self._expand_single_argument(arg2.expand()) + + addends = [a*b*Covariance(*sorted([r1, r2], key=default_sort_key), condition=condition) + for (a, r1) in coeff_rv_list1 for (b, r2) in coeff_rv_list2] + return Add.fromiter(addends) + + @classmethod + def _expand_single_argument(cls, expr): + # return (coefficient, random_symbol) pairs: + if isinstance(expr, RandomSymbol): + return [(S.One, expr)] + elif isinstance(expr, Add): + outval = [] + for a in expr.args: + if isinstance(a, Mul): + outval.append(cls._get_mul_nonrv_rv_tuple(a)) + elif is_random(a): + outval.append((S.One, a)) + + return outval + elif isinstance(expr, Mul): + return [cls._get_mul_nonrv_rv_tuple(expr)] + elif is_random(expr): + return [(S.One, expr)] + + @classmethod + def _get_mul_nonrv_rv_tuple(cls, m): + rv = [] + nonrv = [] + for a in m.args: + if is_random(a): + rv.append(a) + else: + nonrv.append(a) + return (Mul.fromiter(nonrv), Mul.fromiter(rv)) + + def _eval_rewrite_as_Expectation(self, arg1, arg2, condition=None, **kwargs): + e1 = Expectation(arg1*arg2, condition) + e2 = Expectation(arg1, condition)*Expectation(arg2, condition) + return e1 - e2 + + def _eval_rewrite_as_Probability(self, arg1, arg2, condition=None, **kwargs): + return self.rewrite(Expectation).rewrite(Probability) + + def _eval_rewrite_as_Integral(self, arg1, arg2, condition=None, **kwargs): + return covariance(self.args[0], self.args[1], self._condition, evaluate=False) + + _eval_rewrite_as_Sum = _eval_rewrite_as_Integral + + def evaluate_integral(self): + return self.rewrite(Integral).doit() + + +class Moment(Expr): + """ + Symbolic class for Moment + + Examples + ======== + + >>> from sympy import Symbol, Integral + >>> from sympy.stats import Normal, Expectation, Probability, Moment + >>> mu = Symbol('mu', real=True) + >>> sigma = Symbol('sigma', positive=True) + >>> X = Normal('X', mu, sigma) + >>> M = Moment(X, 3, 1) + + To evaluate the result of Moment use `doit`: + + >>> M.doit() + mu**3 - 3*mu**2 + 3*mu*sigma**2 + 3*mu - 3*sigma**2 - 1 + + Rewrite the Moment expression in terms of Expectation: + + >>> M.rewrite(Expectation) + Expectation((X - 1)**3) + + Rewrite the Moment expression in terms of Probability: + + >>> M.rewrite(Probability) + Integral((x - 1)**3*Probability(Eq(X, x)), (x, -oo, oo)) + + Rewrite the Moment expression in terms of Integral: + + >>> M.rewrite(Integral) + Integral(sqrt(2)*(X - 1)**3*exp(-(X - mu)**2/(2*sigma**2))/(2*sqrt(pi)*sigma), (X, -oo, oo)) + + """ + def __new__(cls, X, n, c=0, condition=None, **kwargs): + X = _sympify(X) + n = _sympify(n) + c = _sympify(c) + if condition is not None: + condition = _sympify(condition) + return super().__new__(cls, X, n, c, condition) + else: + return super().__new__(cls, X, n, c) + + def doit(self, **hints): + return self.rewrite(Expectation).doit(**hints) + + def _eval_rewrite_as_Expectation(self, X, n, c=0, condition=None, **kwargs): + return Expectation((X - c)**n, condition) + + def _eval_rewrite_as_Probability(self, X, n, c=0, condition=None, **kwargs): + return self.rewrite(Expectation).rewrite(Probability) + + def _eval_rewrite_as_Integral(self, X, n, c=0, condition=None, **kwargs): + return self.rewrite(Expectation).rewrite(Integral) + + +class CentralMoment(Expr): + """ + Symbolic class Central Moment + + Examples + ======== + + >>> from sympy import Symbol, Integral + >>> from sympy.stats import Normal, Expectation, Probability, CentralMoment + >>> mu = Symbol('mu', real=True) + >>> sigma = Symbol('sigma', positive=True) + >>> X = Normal('X', mu, sigma) + >>> CM = CentralMoment(X, 4) + + To evaluate the result of CentralMoment use `doit`: + + >>> CM.doit().simplify() + 3*sigma**4 + + Rewrite the CentralMoment expression in terms of Expectation: + + >>> CM.rewrite(Expectation) + Expectation((-Expectation(X) + X)**4) + + Rewrite the CentralMoment expression in terms of Probability: + + >>> CM.rewrite(Probability) + Integral((x - Integral(x*Probability(True), (x, -oo, oo)))**4*Probability(Eq(X, x)), (x, -oo, oo)) + + Rewrite the CentralMoment expression in terms of Integral: + + >>> CM.rewrite(Integral) + Integral(sqrt(2)*(X - Integral(sqrt(2)*X*exp(-(X - mu)**2/(2*sigma**2))/(2*sqrt(pi)*sigma), (X, -oo, oo)))**4*exp(-(X - mu)**2/(2*sigma**2))/(2*sqrt(pi)*sigma), (X, -oo, oo)) + + """ + def __new__(cls, X, n, condition=None, **kwargs): + X = _sympify(X) + n = _sympify(n) + if condition is not None: + condition = _sympify(condition) + return super().__new__(cls, X, n, condition) + else: + return super().__new__(cls, X, n) + + def doit(self, **hints): + return self.rewrite(Expectation).doit(**hints) + + def _eval_rewrite_as_Expectation(self, X, n, condition=None, **kwargs): + mu = Expectation(X, condition, **kwargs) + return Moment(X, n, mu, condition, **kwargs).rewrite(Expectation) + + def _eval_rewrite_as_Probability(self, X, n, condition=None, **kwargs): + return self.rewrite(Expectation).rewrite(Probability) + + def _eval_rewrite_as_Integral(self, X, n, condition=None, **kwargs): + return self.rewrite(Expectation).rewrite(Integral) diff --git a/lib/python3.12/site-packages/sympy/stats/tests/__init__.py b/lib/python3.12/site-packages/sympy/stats/tests/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/sympy/stats/tests/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/sympy/stats/tests/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0ab90fea027c290ca3aa89026ce2104bf576d935 Binary files /dev/null and b/lib/python3.12/site-packages/sympy/stats/tests/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sympy/stats/tests/__pycache__/test_compound_rv.cpython-312.pyc b/lib/python3.12/site-packages/sympy/stats/tests/__pycache__/test_compound_rv.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..69c1f3b0efadbc428315b57ecc2d235ec7bec7e6 Binary files /dev/null and b/lib/python3.12/site-packages/sympy/stats/tests/__pycache__/test_compound_rv.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sympy/stats/tests/__pycache__/test_discrete_rv.cpython-312.pyc b/lib/python3.12/site-packages/sympy/stats/tests/__pycache__/test_discrete_rv.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ed3e91de903627748d5cdec5d980557410a8fc89 Binary files /dev/null and b/lib/python3.12/site-packages/sympy/stats/tests/__pycache__/test_discrete_rv.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sympy/stats/tests/__pycache__/test_error_prop.cpython-312.pyc b/lib/python3.12/site-packages/sympy/stats/tests/__pycache__/test_error_prop.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3069cf3e5b612422b9026ee1a0308ba6004acc90 Binary files /dev/null and b/lib/python3.12/site-packages/sympy/stats/tests/__pycache__/test_error_prop.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sympy/stats/tests/__pycache__/test_finite_rv.cpython-312.pyc b/lib/python3.12/site-packages/sympy/stats/tests/__pycache__/test_finite_rv.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..454c207bd2a27dc3de86d62661a229dc97462d80 Binary files /dev/null and b/lib/python3.12/site-packages/sympy/stats/tests/__pycache__/test_finite_rv.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sympy/stats/tests/__pycache__/test_joint_rv.cpython-312.pyc b/lib/python3.12/site-packages/sympy/stats/tests/__pycache__/test_joint_rv.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..04cc3cc753fccc4269e283e4e29ade0deeab4a0f Binary files /dev/null and b/lib/python3.12/site-packages/sympy/stats/tests/__pycache__/test_joint_rv.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sympy/stats/tests/__pycache__/test_matrix_distributions.cpython-312.pyc b/lib/python3.12/site-packages/sympy/stats/tests/__pycache__/test_matrix_distributions.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..deb3dc5be020c9514def7ba1004ae8d7b01868b4 Binary files /dev/null and b/lib/python3.12/site-packages/sympy/stats/tests/__pycache__/test_matrix_distributions.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sympy/stats/tests/__pycache__/test_mix.cpython-312.pyc b/lib/python3.12/site-packages/sympy/stats/tests/__pycache__/test_mix.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..78f7e393631e898bc91acc6ea0184082cc6a37d0 Binary files /dev/null and b/lib/python3.12/site-packages/sympy/stats/tests/__pycache__/test_mix.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sympy/stats/tests/__pycache__/test_random_matrix.cpython-312.pyc b/lib/python3.12/site-packages/sympy/stats/tests/__pycache__/test_random_matrix.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..84879c53ed982e879239a253c3c74c26931cbd1d Binary files /dev/null and b/lib/python3.12/site-packages/sympy/stats/tests/__pycache__/test_random_matrix.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sympy/stats/tests/__pycache__/test_rv.cpython-312.pyc b/lib/python3.12/site-packages/sympy/stats/tests/__pycache__/test_rv.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..edf93b9ac7d27b7367da27f6705e215c5cf773c9 Binary files /dev/null and b/lib/python3.12/site-packages/sympy/stats/tests/__pycache__/test_rv.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sympy/stats/tests/__pycache__/test_stochastic_process.cpython-312.pyc b/lib/python3.12/site-packages/sympy/stats/tests/__pycache__/test_stochastic_process.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9eec33d0ad5c12e8b69efdfd7e7abeb2dfa302b7 Binary files /dev/null and b/lib/python3.12/site-packages/sympy/stats/tests/__pycache__/test_stochastic_process.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sympy/stats/tests/__pycache__/test_symbolic_multivariate.cpython-312.pyc b/lib/python3.12/site-packages/sympy/stats/tests/__pycache__/test_symbolic_multivariate.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..000f26283bc489ba792009262d7da6a43f75ccbf Binary files /dev/null and b/lib/python3.12/site-packages/sympy/stats/tests/__pycache__/test_symbolic_multivariate.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sympy/stats/tests/__pycache__/test_symbolic_probability.cpython-312.pyc b/lib/python3.12/site-packages/sympy/stats/tests/__pycache__/test_symbolic_probability.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1802dc26dcf7b7b1b3f3cc2e78a65fccfdb9d3f9 Binary files /dev/null and b/lib/python3.12/site-packages/sympy/stats/tests/__pycache__/test_symbolic_probability.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sympy/stats/tests/test_compound_rv.py b/lib/python3.12/site-packages/sympy/stats/tests/test_compound_rv.py new file mode 100644 index 0000000000000000000000000000000000000000..573ba364b686738e56bb1c4615acd2a9bc8bf3ae --- /dev/null +++ b/lib/python3.12/site-packages/sympy/stats/tests/test_compound_rv.py @@ -0,0 +1,159 @@ +from sympy.concrete.summations import Sum +from sympy.core.numbers import (oo, pi) +from sympy.core.relational import Eq +from sympy.core.singleton import S +from sympy.core.symbol import symbols +from sympy.functions.combinatorial.factorials import factorial +from sympy.functions.elementary.exponential import exp +from sympy.functions.elementary.miscellaneous import sqrt +from sympy.functions.elementary.piecewise import Piecewise +from sympy.functions.special.beta_functions import beta +from sympy.functions.special.error_functions import erf +from sympy.functions.special.gamma_functions import gamma +from sympy.integrals.integrals import Integral +from sympy.sets.sets import Interval +from sympy.stats import (Normal, P, E, density, Gamma, Poisson, Rayleigh, + variance, Bernoulli, Beta, Uniform, cdf) +from sympy.stats.compound_rv import CompoundDistribution, CompoundPSpace +from sympy.stats.crv_types import NormalDistribution +from sympy.stats.drv_types import PoissonDistribution +from sympy.stats.frv_types import BernoulliDistribution +from sympy.testing.pytest import raises, ignore_warnings +from sympy.stats.joint_rv_types import MultivariateNormalDistribution + +from sympy.abc import x + + +# helpers for testing troublesome unevaluated expressions +flat = lambda s: ''.join(str(s).split()) +streq = lambda *a: len(set(map(flat, a))) == 1 +assert streq(x, x) +assert streq(x, 'x') +assert not streq(x, x + 1) + + +def test_normal_CompoundDist(): + X = Normal('X', 1, 2) + Y = Normal('X', X, 4) + assert density(Y)(x).simplify() == sqrt(10)*exp(-x**2/40 + x/20 - S(1)/40)/(20*sqrt(pi)) + assert E(Y) == 1 # it is always equal to mean of X + assert P(Y > 1) == S(1)/2 # as 1 is the mean + assert P(Y > 5).simplify() == S(1)/2 - erf(sqrt(10)/5)/2 + assert variance(Y) == variance(X) + 4**2 # 2**2 + 4**2 + # https://math.stackexchange.com/questions/1484451/ + # (Contains proof of E and variance computation) + + +def test_poisson_CompoundDist(): + k, t, y = symbols('k t y', positive=True, real=True) + G = Gamma('G', k, t) + D = Poisson('P', G) + assert density(D)(y).simplify() == t**y*(t + 1)**(-k - y)*gamma(k + y)/(gamma(k)*gamma(y + 1)) + # https://en.wikipedia.org/wiki/Negative_binomial_distribution#Gamma%E2%80%93Poisson_mixture + assert E(D).simplify() == k*t # mean of NegativeBinomialDistribution + + +def test_bernoulli_CompoundDist(): + X = Beta('X', 1, 2) + Y = Bernoulli('Y', X) + assert density(Y).dict == {0: S(2)/3, 1: S(1)/3} + assert E(Y) == P(Eq(Y, 1)) == S(1)/3 + assert variance(Y) == S(2)/9 + assert cdf(Y) == {0: S(2)/3, 1: 1} + + # test issue 8128 + a = Bernoulli('a', S(1)/2) + b = Bernoulli('b', a) + assert density(b).dict == {0: S(1)/2, 1: S(1)/2} + assert P(b > 0.5) == S(1)/2 + + X = Uniform('X', 0, 1) + Y = Bernoulli('Y', X) + assert E(Y) == S(1)/2 + assert P(Eq(Y, 1)) == E(Y) + + +def test_unevaluated_CompoundDist(): + # these tests need to be removed once they work with evaluation as they are currently not + # evaluated completely in sympy. + R = Rayleigh('R', 4) + X = Normal('X', 3, R) + ans = ''' + Piecewise(((-sqrt(pi)*sinh(x/4 - 3/4) + sqrt(pi)*cosh(x/4 - 3/4))/( + 8*sqrt(pi)), Abs(arg(x - 3)) <= pi/4), (Integral(sqrt(2)*exp(-(x - 3) + **2/(2*R**2))*exp(-R**2/32)/(32*sqrt(pi)), (R, 0, oo)), True))''' + assert streq(density(X)(x), ans) + + expre = ''' + Integral(X*Integral(sqrt(2)*exp(-(X-3)**2/(2*R**2))*exp(-R**2/32)/(32* + sqrt(pi)),(R,0,oo)),(X,-oo,oo))''' + with ignore_warnings(UserWarning): ### TODO: Restore tests once warnings are removed + assert streq(E(X, evaluate=False).rewrite(Integral), expre) + + X = Poisson('X', 1) + Y = Poisson('Y', X) + Z = Poisson('Z', Y) + exprd = Sum(exp(-Y)*Y**x*Sum(exp(-1)*exp(-X)*X**Y/(factorial(X)*factorial(Y) + ), (X, 0, oo))/factorial(x), (Y, 0, oo)) + assert density(Z)(x) == exprd + + N = Normal('N', 1, 2) + M = Normal('M', 3, 4) + D = Normal('D', M, N) + exprd = ''' + Integral(sqrt(2)*exp(-(N-1)**2/8)*Integral(exp(-(x-M)**2/(2*N**2))*exp + (-(M-3)**2/32)/(8*pi*N),(M,-oo,oo))/(4*sqrt(pi)),(N,-oo,oo))''' + assert streq(density(D, evaluate=False)(x), exprd) + + +def test_Compound_Distribution(): + X = Normal('X', 2, 4) + N = NormalDistribution(X, 4) + C = CompoundDistribution(N) + assert C.is_Continuous + assert C.set == Interval(-oo, oo) + assert C.pdf(x, evaluate=True).simplify() == exp(-x**2/64 + x/16 - S(1)/16)/(8*sqrt(pi)) + + assert not isinstance(CompoundDistribution(NormalDistribution(2, 3)), + CompoundDistribution) + M = MultivariateNormalDistribution([1, 2], [[2, 1], [1, 2]]) + raises(NotImplementedError, lambda: CompoundDistribution(M)) + + X = Beta('X', 2, 4) + B = BernoulliDistribution(X, 1, 0) + C = CompoundDistribution(B) + assert C.is_Finite + assert C.set == {0, 1} + y = symbols('y', negative=False, integer=True) + assert C.pdf(y, evaluate=True) == Piecewise((S(1)/(30*beta(2, 4)), Eq(y, 0)), + (S(1)/(60*beta(2, 4)), Eq(y, 1)), (0, True)) + + k, t, z = symbols('k t z', positive=True, real=True) + G = Gamma('G', k, t) + X = PoissonDistribution(G) + C = CompoundDistribution(X) + assert C.is_Discrete + assert C.set == S.Naturals0 + assert C.pdf(z, evaluate=True).simplify() == t**z*(t + 1)**(-k - z)*gamma(k \ + + z)/(gamma(k)*gamma(z + 1)) + + +def test_compound_pspace(): + X = Normal('X', 2, 4) + Y = Normal('Y', 3, 6) + assert not isinstance(Y.pspace, CompoundPSpace) + N = NormalDistribution(1, 2) + D = PoissonDistribution(3) + B = BernoulliDistribution(0.2, 1, 0) + pspace1 = CompoundPSpace('N', N) + pspace2 = CompoundPSpace('D', D) + pspace3 = CompoundPSpace('B', B) + assert not isinstance(pspace1, CompoundPSpace) + assert not isinstance(pspace2, CompoundPSpace) + assert not isinstance(pspace3, CompoundPSpace) + M = MultivariateNormalDistribution([1, 2], [[2, 1], [1, 2]]) + raises(ValueError, lambda: CompoundPSpace('M', M)) + Y = Normal('Y', X, 6) + assert isinstance(Y.pspace, CompoundPSpace) + assert Y.pspace.distribution == CompoundDistribution(NormalDistribution(X, 6)) + assert Y.pspace.domain.set == Interval(-oo, oo) diff --git a/lib/python3.12/site-packages/sympy/stats/tests/test_continuous_rv.py b/lib/python3.12/site-packages/sympy/stats/tests/test_continuous_rv.py new file mode 100644 index 0000000000000000000000000000000000000000..b2c4206b5c29ffd3194d1ae05e57c51c9c1b6d78 --- /dev/null +++ b/lib/python3.12/site-packages/sympy/stats/tests/test_continuous_rv.py @@ -0,0 +1,1583 @@ +from sympy.concrete.summations import Sum +from sympy.core.function import (Lambda, diff, expand_func) +from sympy.core.mul import Mul +from sympy.core import EulerGamma +from sympy.core.numbers import (E as e, I, Rational, pi) +from sympy.core.relational import (Eq, Ne) +from sympy.core.singleton import S +from sympy.core.symbol import (Dummy, Symbol, symbols) +from sympy.functions.combinatorial.factorials import (binomial, factorial) +from sympy.functions.elementary.complexes import (Abs, im, re, sign) +from sympy.functions.elementary.exponential import (exp, log) +from sympy.functions.elementary.hyperbolic import (cosh, sinh) +from sympy.functions.elementary.integers import floor +from sympy.functions.elementary.miscellaneous import sqrt +from sympy.functions.elementary.piecewise import Piecewise +from sympy.functions.elementary.trigonometric import (asin, atan, cos, sin, tan) +from sympy.functions.special.bessel import (besseli, besselj, besselk) +from sympy.functions.special.beta_functions import beta +from sympy.functions.special.error_functions import (erf, erfc, erfi, expint) +from sympy.functions.special.gamma_functions import (gamma, lowergamma, uppergamma) +from sympy.functions.special.zeta_functions import zeta +from sympy.functions.special.hyper import hyper +from sympy.integrals.integrals import Integral +from sympy.logic.boolalg import (And, Or) +from sympy.sets.sets import Interval +from sympy.simplify.simplify import simplify +from sympy.utilities.lambdify import lambdify +from sympy.functions.special.error_functions import erfinv +from sympy.functions.special.hyper import meijerg +from sympy.sets.sets import FiniteSet, Complement, Intersection +from sympy.stats import (P, E, where, density, variance, covariance, skewness, kurtosis, median, + given, pspace, cdf, characteristic_function, moment_generating_function, + ContinuousRV, Arcsin, Benini, Beta, BetaNoncentral, BetaPrime, + Cauchy, Chi, ChiSquared, ChiNoncentral, Dagum, Davis, Erlang, ExGaussian, + Exponential, ExponentialPower, FDistribution, FisherZ, Frechet, Gamma, + GammaInverse, Gompertz, Gumbel, Kumaraswamy, Laplace, Levy, Logistic, LogCauchy, + LogLogistic, LogitNormal, LogNormal, Maxwell, Moyal, Nakagami, Normal, GaussianInverse, + Pareto, PowerFunction, QuadraticU, RaisedCosine, Rayleigh, Reciprocal, ShiftedGompertz, StudentT, + Trapezoidal, Triangular, Uniform, UniformSum, VonMises, Weibull, coskewness, + WignerSemicircle, Wald, correlation, moment, cmoment, smoment, quantile, + Lomax, BoundedPareto) + +from sympy.stats.crv_types import NormalDistribution, ExponentialDistribution, ContinuousDistributionHandmade +from sympy.stats.joint_rv_types import MultivariateLaplaceDistribution, MultivariateNormalDistribution +from sympy.stats.crv import SingleContinuousPSpace, SingleContinuousDomain +from sympy.stats.compound_rv import CompoundPSpace +from sympy.stats.symbolic_probability import Probability +from sympy.testing.pytest import raises, XFAIL, slow, ignore_warnings +from sympy.core.random import verify_numerically as tn + +oo = S.Infinity + +x, y, z = map(Symbol, 'xyz') + +def test_single_normal(): + mu = Symbol('mu', real=True) + sigma = Symbol('sigma', positive=True) + X = Normal('x', 0, 1) + Y = X*sigma + mu + + assert E(Y) == mu + assert variance(Y) == sigma**2 + pdf = density(Y) + x = Symbol('x', real=True) + assert (pdf(x) == + 2**S.Half*exp(-(x - mu)**2/(2*sigma**2))/(2*pi**S.Half*sigma)) + + assert P(X**2 < 1) == erf(2**S.Half/2) + ans = quantile(Y)(x) + assert ans == Complement(Intersection(FiniteSet( + sqrt(2)*sigma*(sqrt(2)*mu/(2*sigma)+ erfinv(2*x - 1))), + Interval(-oo, oo)), FiniteSet(mu)) + assert E(X, Eq(X, mu)) == mu + + assert median(X) == FiniteSet(0) + # issue 8248 + assert X.pspace.compute_expectation(1).doit() == 1 + + +def test_conditional_1d(): + X = Normal('x', 0, 1) + Y = given(X, X >= 0) + z = Symbol('z') + + assert density(Y)(z) == 2 * density(X)(z) + + assert Y.pspace.domain.set == Interval(0, oo) + assert E(Y) == sqrt(2) / sqrt(pi) + + assert E(X**2) == E(Y**2) + + +def test_ContinuousDomain(): + X = Normal('x', 0, 1) + assert where(X**2 <= 1).set == Interval(-1, 1) + assert where(X**2 <= 1).symbol == X.symbol + assert where(And(X**2 <= 1, X >= 0)).set == Interval(0, 1) + raises(ValueError, lambda: where(sin(X) > 1)) + + Y = given(X, X >= 0) + + assert Y.pspace.domain.set == Interval(0, oo) + + +def test_multiple_normal(): + X, Y = Normal('x', 0, 1), Normal('y', 0, 1) + p = Symbol("p", positive=True) + + assert E(X + Y) == 0 + assert variance(X + Y) == 2 + assert variance(X + X) == 4 + assert covariance(X, Y) == 0 + assert covariance(2*X + Y, -X) == -2*variance(X) + assert skewness(X) == 0 + assert skewness(X + Y) == 0 + assert kurtosis(X) == 3 + assert kurtosis(X+Y) == 3 + assert correlation(X, Y) == 0 + assert correlation(X, X + Y) == correlation(X, X - Y) + assert moment(X, 2) == 1 + assert cmoment(X, 3) == 0 + assert moment(X + Y, 4) == 12 + assert cmoment(X, 2) == variance(X) + assert smoment(X*X, 2) == 1 + assert smoment(X + Y, 3) == skewness(X + Y) + assert smoment(X + Y, 4) == kurtosis(X + Y) + assert E(X, Eq(X + Y, 0)) == 0 + assert variance(X, Eq(X + Y, 0)) == S.Half + assert quantile(X)(p) == sqrt(2)*erfinv(2*p - S.One) + + +def test_symbolic(): + mu1, mu2 = symbols('mu1 mu2', real=True) + s1, s2 = symbols('sigma1 sigma2', positive=True) + rate = Symbol('lambda', positive=True) + X = Normal('x', mu1, s1) + Y = Normal('y', mu2, s2) + Z = Exponential('z', rate) + a, b, c = symbols('a b c', real=True) + + assert E(X) == mu1 + assert E(X + Y) == mu1 + mu2 + assert E(a*X + b) == a*E(X) + b + assert variance(X) == s1**2 + assert variance(X + a*Y + b) == variance(X) + a**2*variance(Y) + + assert E(Z) == 1/rate + assert E(a*Z + b) == a*E(Z) + b + assert E(X + a*Z + b) == mu1 + a/rate + b + assert median(X) == FiniteSet(mu1) + + +def test_cdf(): + X = Normal('x', 0, 1) + + d = cdf(X) + assert P(X < 1) == d(1).rewrite(erfc) + assert d(0) == S.Half + + d = cdf(X, X > 0) # given X>0 + assert d(0) == 0 + + Y = Exponential('y', 10) + d = cdf(Y) + assert d(-5) == 0 + assert P(Y > 3) == 1 - d(3) + + raises(ValueError, lambda: cdf(X + Y)) + + Z = Exponential('z', 1) + f = cdf(Z) + assert f(z) == Piecewise((1 - exp(-z), z >= 0), (0, True)) + + +def test_characteristic_function(): + X = Uniform('x', 0, 1) + + cf = characteristic_function(X) + assert cf(1) == -I*(-1 + exp(I)) + + Y = Normal('y', 1, 1) + cf = characteristic_function(Y) + assert cf(0) == 1 + assert cf(1) == exp(I - S.Half) + + Z = Exponential('z', 5) + cf = characteristic_function(Z) + assert cf(0) == 1 + assert cf(1).expand() == Rational(25, 26) + I*5/26 + + X = GaussianInverse('x', 1, 1) + cf = characteristic_function(X) + assert cf(0) == 1 + assert cf(1) == exp(1 - sqrt(1 - 2*I)) + + X = ExGaussian('x', 0, 1, 1) + cf = characteristic_function(X) + assert cf(0) == 1 + assert cf(1) == (1 + I)*exp(Rational(-1, 2))/2 + + L = Levy('x', 0, 1) + cf = characteristic_function(L) + assert cf(0) == 1 + assert cf(1) == exp(-sqrt(2)*sqrt(-I)) + + +def test_moment_generating_function(): + t = symbols('t', positive=True) + + # Symbolic tests + a, b, c = symbols('a b c') + + mgf = moment_generating_function(Beta('x', a, b))(t) + assert mgf == hyper((a,), (a + b,), t) + + mgf = moment_generating_function(Chi('x', a))(t) + assert mgf == sqrt(2)*t*gamma(a/2 + S.Half)*\ + hyper((a/2 + S.Half,), (Rational(3, 2),), t**2/2)/gamma(a/2) +\ + hyper((a/2,), (S.Half,), t**2/2) + + mgf = moment_generating_function(ChiSquared('x', a))(t) + assert mgf == (1 - 2*t)**(-a/2) + + mgf = moment_generating_function(Erlang('x', a, b))(t) + assert mgf == (1 - t/b)**(-a) + + mgf = moment_generating_function(ExGaussian("x", a, b, c))(t) + assert mgf == exp(a*t + b**2*t**2/2)/(1 - t/c) + + mgf = moment_generating_function(Exponential('x', a))(t) + assert mgf == a/(a - t) + + mgf = moment_generating_function(Gamma('x', a, b))(t) + assert mgf == (-b*t + 1)**(-a) + + mgf = moment_generating_function(Gumbel('x', a, b))(t) + assert mgf == exp(b*t)*gamma(-a*t + 1) + + mgf = moment_generating_function(Gompertz('x', a, b))(t) + assert mgf == b*exp(b)*expint(t/a, b) + + mgf = moment_generating_function(Laplace('x', a, b))(t) + assert mgf == exp(a*t)/(-b**2*t**2 + 1) + + mgf = moment_generating_function(Logistic('x', a, b))(t) + assert mgf == exp(a*t)*beta(-b*t + 1, b*t + 1) + + mgf = moment_generating_function(Normal('x', a, b))(t) + assert mgf == exp(a*t + b**2*t**2/2) + + mgf = moment_generating_function(Pareto('x', a, b))(t) + assert mgf == b*(-a*t)**b*uppergamma(-b, -a*t) + + mgf = moment_generating_function(QuadraticU('x', a, b))(t) + assert str(mgf) == ("(3*(t*(-4*b + (a + b)**2) + 4)*exp(b*t) - " + "3*(t*(a**2 + 2*a*(b - 2) + b**2) + 4)*exp(a*t))/(t**2*(a - b)**3)") + + mgf = moment_generating_function(RaisedCosine('x', a, b))(t) + assert mgf == pi**2*exp(a*t)*sinh(b*t)/(b*t*(b**2*t**2 + pi**2)) + + mgf = moment_generating_function(Rayleigh('x', a))(t) + assert mgf == sqrt(2)*sqrt(pi)*a*t*(erf(sqrt(2)*a*t/2) + 1)\ + *exp(a**2*t**2/2)/2 + 1 + + mgf = moment_generating_function(Triangular('x', a, b, c))(t) + assert str(mgf) == ("(-2*(-a + b)*exp(c*t) + 2*(-a + c)*exp(b*t) + " + "2*(b - c)*exp(a*t))/(t**2*(-a + b)*(-a + c)*(b - c))") + + mgf = moment_generating_function(Uniform('x', a, b))(t) + assert mgf == (-exp(a*t) + exp(b*t))/(t*(-a + b)) + + mgf = moment_generating_function(UniformSum('x', a))(t) + assert mgf == ((exp(t) - 1)/t)**a + + mgf = moment_generating_function(WignerSemicircle('x', a))(t) + assert mgf == 2*besseli(1, a*t)/(a*t) + + # Numeric tests + + mgf = moment_generating_function(Beta('x', 1, 1))(t) + assert mgf.diff(t).subs(t, 1) == hyper((2,), (3,), 1)/2 + + mgf = moment_generating_function(Chi('x', 1))(t) + assert mgf.diff(t).subs(t, 1) == sqrt(2)*hyper((1,), (Rational(3, 2),), S.Half + )/sqrt(pi) + hyper((Rational(3, 2),), (Rational(3, 2),), S.Half) + 2*sqrt(2)*hyper((2,), + (Rational(5, 2),), S.Half)/(3*sqrt(pi)) + + mgf = moment_generating_function(ChiSquared('x', 1))(t) + assert mgf.diff(t).subs(t, 1) == I + + mgf = moment_generating_function(Erlang('x', 1, 1))(t) + assert mgf.diff(t).subs(t, 0) == 1 + + mgf = moment_generating_function(ExGaussian("x", 0, 1, 1))(t) + assert mgf.diff(t).subs(t, 2) == -exp(2) + + mgf = moment_generating_function(Exponential('x', 1))(t) + assert mgf.diff(t).subs(t, 0) == 1 + + mgf = moment_generating_function(Gamma('x', 1, 1))(t) + assert mgf.diff(t).subs(t, 0) == 1 + + mgf = moment_generating_function(Gumbel('x', 1, 1))(t) + assert mgf.diff(t).subs(t, 0) == EulerGamma + 1 + + mgf = moment_generating_function(Gompertz('x', 1, 1))(t) + assert mgf.diff(t).subs(t, 1) == -e*meijerg(((), (1, 1)), + ((0, 0, 0), ()), 1) + + mgf = moment_generating_function(Laplace('x', 1, 1))(t) + assert mgf.diff(t).subs(t, 0) == 1 + + mgf = moment_generating_function(Logistic('x', 1, 1))(t) + assert mgf.diff(t).subs(t, 0) == beta(1, 1) + + mgf = moment_generating_function(Normal('x', 0, 1))(t) + assert mgf.diff(t).subs(t, 1) == exp(S.Half) + + mgf = moment_generating_function(Pareto('x', 1, 1))(t) + assert mgf.diff(t).subs(t, 0) == expint(1, 0) + + mgf = moment_generating_function(QuadraticU('x', 1, 2))(t) + assert mgf.diff(t).subs(t, 1) == -12*e - 3*exp(2) + + mgf = moment_generating_function(RaisedCosine('x', 1, 1))(t) + assert mgf.diff(t).subs(t, 1) == -2*e*pi**2*sinh(1)/\ + (1 + pi**2)**2 + e*pi**2*cosh(1)/(1 + pi**2) + + mgf = moment_generating_function(Rayleigh('x', 1))(t) + assert mgf.diff(t).subs(t, 0) == sqrt(2)*sqrt(pi)/2 + + mgf = moment_generating_function(Triangular('x', 1, 3, 2))(t) + assert mgf.diff(t).subs(t, 1) == -e + exp(3) + + mgf = moment_generating_function(Uniform('x', 0, 1))(t) + assert mgf.diff(t).subs(t, 1) == 1 + + mgf = moment_generating_function(UniformSum('x', 1))(t) + assert mgf.diff(t).subs(t, 1) == 1 + + mgf = moment_generating_function(WignerSemicircle('x', 1))(t) + assert mgf.diff(t).subs(t, 1) == -2*besseli(1, 1) + besseli(2, 1) +\ + besseli(0, 1) + + +def test_ContinuousRV(): + pdf = sqrt(2)*exp(-x**2/2)/(2*sqrt(pi)) # Normal distribution + # X and Y should be equivalent + X = ContinuousRV(x, pdf, check=True) + Y = Normal('y', 0, 1) + + assert variance(X) == variance(Y) + assert P(X > 0) == P(Y > 0) + Z = ContinuousRV(z, exp(-z), set=Interval(0, oo)) + assert Z.pspace.domain.set == Interval(0, oo) + assert E(Z) == 1 + assert P(Z > 5) == exp(-5) + raises(ValueError, lambda: ContinuousRV(z, exp(-z), set=Interval(0, 10), check=True)) + + # the correct pdf for Gamma(k, theta) but the integral in `check` + # integrates to something equivalent to 1 and not to 1 exactly + _x, k, theta = symbols("x k theta", positive=True) + pdf = 1/(gamma(k)*theta**k)*_x**(k-1)*exp(-_x/theta) + X = ContinuousRV(_x, pdf, set=Interval(0, oo)) + Y = Gamma('y', k, theta) + assert (E(X) - E(Y)).simplify() == 0 + assert (variance(X) - variance(Y)).simplify() == 0 + + +def test_arcsin(): + + a = Symbol("a", real=True) + b = Symbol("b", real=True) + + X = Arcsin('x', a, b) + assert density(X)(x) == 1/(pi*sqrt((-x + b)*(x - a))) + assert cdf(X)(x) == Piecewise((0, a > x), + (2*asin(sqrt((-a + x)/(-a + b)))/pi, b >= x), + (1, True)) + assert pspace(X).domain.set == Interval(a, b) + +def test_benini(): + alpha = Symbol("alpha", positive=True) + beta = Symbol("beta", positive=True) + sigma = Symbol("sigma", positive=True) + X = Benini('x', alpha, beta, sigma) + + assert density(X)(x) == ((alpha/x + 2*beta*log(x/sigma)/x) + *exp(-alpha*log(x/sigma) - beta*log(x/sigma)**2)) + + assert pspace(X).domain.set == Interval(sigma, oo) + raises(NotImplementedError, lambda: moment_generating_function(X)) + alpha = Symbol("alpha", nonpositive=True) + raises(ValueError, lambda: Benini('x', alpha, beta, sigma)) + + beta = Symbol("beta", nonpositive=True) + raises(ValueError, lambda: Benini('x', alpha, beta, sigma)) + + alpha = Symbol("alpha", positive=True) + raises(ValueError, lambda: Benini('x', alpha, beta, sigma)) + + beta = Symbol("beta", positive=True) + sigma = Symbol("sigma", nonpositive=True) + raises(ValueError, lambda: Benini('x', alpha, beta, sigma)) + +def test_beta(): + a, b = symbols('alpha beta', positive=True) + B = Beta('x', a, b) + + assert pspace(B).domain.set == Interval(0, 1) + assert characteristic_function(B)(x) == hyper((a,), (a + b,), I*x) + assert density(B)(x) == x**(a - 1)*(1 - x)**(b - 1)/beta(a, b) + + assert simplify(E(B)) == a / (a + b) + assert simplify(variance(B)) == a*b / (a**3 + 3*a**2*b + a**2 + 3*a*b**2 + 2*a*b + b**3 + b**2) + + # Full symbolic solution is too much, test with numeric version + a, b = 1, 2 + B = Beta('x', a, b) + assert expand_func(E(B)) == a / S(a + b) + assert expand_func(variance(B)) == (a*b) / S((a + b)**2 * (a + b + 1)) + assert median(B) == FiniteSet(1 - 1/sqrt(2)) + +def test_beta_noncentral(): + a, b = symbols('a b', positive=True) + c = Symbol('c', nonnegative=True) + _k = Dummy('k') + + X = BetaNoncentral('x', a, b, c) + + assert pspace(X).domain.set == Interval(0, 1) + + dens = density(X) + z = Symbol('z') + + res = Sum( z**(_k + a - 1)*(c/2)**_k*(1 - z)**(b - 1)*exp(-c/2)/ + (beta(_k + a, b)*factorial(_k)), (_k, 0, oo)) + assert dens(z).dummy_eq(res) + + # BetaCentral should not raise if the assumptions + # on the symbols can not be determined + a, b, c = symbols('a b c') + assert BetaNoncentral('x', a, b, c) + + a = Symbol('a', positive=False, real=True) + raises(ValueError, lambda: BetaNoncentral('x', a, b, c)) + + a = Symbol('a', positive=True) + b = Symbol('b', positive=False, real=True) + raises(ValueError, lambda: BetaNoncentral('x', a, b, c)) + + a = Symbol('a', positive=True) + b = Symbol('b', positive=True) + c = Symbol('c', nonnegative=False, real=True) + raises(ValueError, lambda: BetaNoncentral('x', a, b, c)) + +def test_betaprime(): + alpha = Symbol("alpha", positive=True) + + betap = Symbol("beta", positive=True) + + X = BetaPrime('x', alpha, betap) + assert density(X)(x) == x**(alpha - 1)*(x + 1)**(-alpha - betap)/beta(alpha, betap) + + alpha = Symbol("alpha", nonpositive=True) + raises(ValueError, lambda: BetaPrime('x', alpha, betap)) + + alpha = Symbol("alpha", positive=True) + betap = Symbol("beta", nonpositive=True) + raises(ValueError, lambda: BetaPrime('x', alpha, betap)) + X = BetaPrime('x', 1, 1) + assert median(X) == FiniteSet(1) + + +def test_BoundedPareto(): + L, H = symbols('L, H', negative=True) + raises(ValueError, lambda: BoundedPareto('X', 1, L, H)) + L, H = symbols('L, H', real=False) + raises(ValueError, lambda: BoundedPareto('X', 1, L, H)) + L, H = symbols('L, H', positive=True) + raises(ValueError, lambda: BoundedPareto('X', -1, L, H)) + + X = BoundedPareto('X', 2, L, H) + assert X.pspace.domain.set == Interval(L, H) + assert density(X)(x) == 2*L**2/(x**3*(1 - L**2/H**2)) + assert cdf(X)(x) == Piecewise((-H**2*L**2/(x**2*(H**2 - L**2)) \ + + H**2/(H**2 - L**2), L <= x), (0, True)) + assert E(X).simplify() == 2*H*L/(H + L) + X = BoundedPareto('X', 1, 2, 4) + assert E(X).simplify() == log(16) + assert median(X) == FiniteSet(Rational(8, 3)) + assert variance(X).simplify() == 8 - 16*log(2)**2 + + +def test_cauchy(): + x0 = Symbol("x0", real=True) + gamma = Symbol("gamma", positive=True) + p = Symbol("p", positive=True) + + X = Cauchy('x', x0, gamma) + # Tests the characteristic function + assert characteristic_function(X)(x) == exp(-gamma*Abs(x) + I*x*x0) + raises(NotImplementedError, lambda: moment_generating_function(X)) + assert density(X)(x) == 1/(pi*gamma*(1 + (x - x0)**2/gamma**2)) + assert diff(cdf(X)(x), x) == density(X)(x) + assert quantile(X)(p) == gamma*tan(pi*(p - S.Half)) + x0 + + x1 = Symbol("x1", real=False) + raises(ValueError, lambda: Cauchy('x', x1, gamma)) + gamma = Symbol("gamma", nonpositive=True) + raises(ValueError, lambda: Cauchy('x', x0, gamma)) + assert median(X) == FiniteSet(x0) + +def test_chi(): + from sympy.core.numbers import I + k = Symbol("k", integer=True) + + X = Chi('x', k) + assert density(X)(x) == 2**(-k/2 + 1)*x**(k - 1)*exp(-x**2/2)/gamma(k/2) + + # Tests the characteristic function + assert characteristic_function(X)(x) == sqrt(2)*I*x*gamma(k/2 + S(1)/2)*hyper((k/2 + S(1)/2,), + (S(3)/2,), -x**2/2)/gamma(k/2) + hyper((k/2,), (S(1)/2,), -x**2/2) + + # Tests the moment generating function + assert moment_generating_function(X)(x) == sqrt(2)*x*gamma(k/2 + S(1)/2)*hyper((k/2 + S(1)/2,), + (S(3)/2,), x**2/2)/gamma(k/2) + hyper((k/2,), (S(1)/2,), x**2/2) + + k = Symbol("k", integer=True, positive=False) + raises(ValueError, lambda: Chi('x', k)) + + k = Symbol("k", integer=False, positive=True) + raises(ValueError, lambda: Chi('x', k)) + +def test_chi_noncentral(): + k = Symbol("k", integer=True) + l = Symbol("l") + + X = ChiNoncentral("x", k, l) + assert density(X)(x) == (x**k*l*(x*l)**(-k/2)* + exp(-x**2/2 - l**2/2)*besseli(k/2 - 1, x*l)) + + k = Symbol("k", integer=True, positive=False) + raises(ValueError, lambda: ChiNoncentral('x', k, l)) + + k = Symbol("k", integer=True, positive=True) + l = Symbol("l", nonpositive=True) + raises(ValueError, lambda: ChiNoncentral('x', k, l)) + + k = Symbol("k", integer=False) + l = Symbol("l", positive=True) + raises(ValueError, lambda: ChiNoncentral('x', k, l)) + + +def test_chi_squared(): + k = Symbol("k", integer=True) + X = ChiSquared('x', k) + + # Tests the characteristic function + assert characteristic_function(X)(x) == ((-2*I*x + 1)**(-k/2)) + + assert density(X)(x) == 2**(-k/2)*x**(k/2 - 1)*exp(-x/2)/gamma(k/2) + assert cdf(X)(x) == Piecewise((lowergamma(k/2, x/2)/gamma(k/2), x >= 0), (0, True)) + assert E(X) == k + assert variance(X) == 2*k + + X = ChiSquared('x', 15) + assert cdf(X)(3) == -14873*sqrt(6)*exp(Rational(-3, 2))/(5005*sqrt(pi)) + erf(sqrt(6)/2) + + k = Symbol("k", integer=True, positive=False) + raises(ValueError, lambda: ChiSquared('x', k)) + + k = Symbol("k", integer=False, positive=True) + raises(ValueError, lambda: ChiSquared('x', k)) + + +def test_dagum(): + p = Symbol("p", positive=True) + b = Symbol("b", positive=True) + a = Symbol("a", positive=True) + + X = Dagum('x', p, a, b) + assert density(X)(x) == a*p*(x/b)**(a*p)*((x/b)**a + 1)**(-p - 1)/x + assert cdf(X)(x) == Piecewise(((1 + (x/b)**(-a))**(-p), x >= 0), + (0, True)) + + p = Symbol("p", nonpositive=True) + raises(ValueError, lambda: Dagum('x', p, a, b)) + + p = Symbol("p", positive=True) + b = Symbol("b", nonpositive=True) + raises(ValueError, lambda: Dagum('x', p, a, b)) + + b = Symbol("b", positive=True) + a = Symbol("a", nonpositive=True) + raises(ValueError, lambda: Dagum('x', p, a, b)) + X = Dagum('x', 1, 1, 1) + assert median(X) == FiniteSet(1) + +def test_davis(): + b = Symbol("b", positive=True) + n = Symbol("n", positive=True) + mu = Symbol("mu", positive=True) + + X = Davis('x', b, n, mu) + dividend = b**n*(x - mu)**(-1-n) + divisor = (exp(b/(x-mu))-1)*(gamma(n)*zeta(n)) + assert density(X)(x) == dividend/divisor + + +def test_erlang(): + k = Symbol("k", integer=True, positive=True) + l = Symbol("l", positive=True) + + X = Erlang("x", k, l) + assert density(X)(x) == x**(k - 1)*l**k*exp(-x*l)/gamma(k) + assert cdf(X)(x) == Piecewise((lowergamma(k, l*x)/gamma(k), x > 0), + (0, True)) + + +def test_exgaussian(): + m, z = symbols("m, z") + s, l = symbols("s, l", positive=True) + X = ExGaussian("x", m, s, l) + + assert density(X)(z) == l*exp(l*(l*s**2 + 2*m - 2*z)/2) *\ + erfc(sqrt(2)*(l*s**2 + m - z)/(2*s))/2 + + # Note: actual_output simplifies to expected_output. + # Ideally cdf(X)(z) would return expected_output + # expected_output = (erf(sqrt(2)*(l*s**2 + m - z)/(2*s)) - 1)*exp(l*(l*s**2 + 2*m - 2*z)/2)/2 - erf(sqrt(2)*(m - z)/(2*s))/2 + S.Half + u = l*(z - m) + v = l*s + GaussianCDF1 = cdf(Normal('x', 0, v))(u) + GaussianCDF2 = cdf(Normal('x', v**2, v))(u) + actual_output = GaussianCDF1 - exp(-u + (v**2/2) + log(GaussianCDF2)) + assert cdf(X)(z) == actual_output + # assert simplify(actual_output) == expected_output + + assert variance(X).expand() == s**2 + l**(-2) + + assert skewness(X).expand() == 2/(l**3*s**2*sqrt(s**2 + l**(-2)) + l * + sqrt(s**2 + l**(-2))) + + +@slow +def test_exponential(): + rate = Symbol('lambda', positive=True) + X = Exponential('x', rate) + p = Symbol("p", positive=True, real=True) + + assert E(X) == 1/rate + assert variance(X) == 1/rate**2 + assert skewness(X) == 2 + assert skewness(X) == smoment(X, 3) + assert kurtosis(X) == 9 + assert kurtosis(X) == smoment(X, 4) + assert smoment(2*X, 4) == smoment(X, 4) + assert moment(X, 3) == 3*2*1/rate**3 + assert P(X > 0) is S.One + assert P(X > 1) == exp(-rate) + assert P(X > 10) == exp(-10*rate) + assert quantile(X)(p) == -log(1-p)/rate + + assert where(X <= 1).set == Interval(0, 1) + Y = Exponential('y', 1) + assert median(Y) == FiniteSet(log(2)) + #Test issue 9970 + z = Dummy('z') + assert P(X > z) == exp(-z*rate) + assert P(X < z) == 0 + #Test issue 10076 (Distribution with interval(0,oo)) + x = Symbol('x') + _z = Dummy('_z') + b = SingleContinuousPSpace(x, ExponentialDistribution(2)) + + with ignore_warnings(UserWarning): ### TODO: Restore tests once warnings are removed + expected1 = Integral(2*exp(-2*_z), (_z, 3, oo)) + assert b.probability(x > 3, evaluate=False).rewrite(Integral).dummy_eq(expected1) + + expected2 = Integral(2*exp(-2*_z), (_z, 0, 4)) + assert b.probability(x < 4, evaluate=False).rewrite(Integral).dummy_eq(expected2) + Y = Exponential('y', 2*rate) + assert coskewness(X, X, X) == skewness(X) + assert coskewness(X, Y + rate*X, Y + 2*rate*X) == \ + 4/(sqrt(1 + 1/(4*rate**2))*sqrt(4 + 1/(4*rate**2))) + assert coskewness(X + 2*Y, Y + X, Y + 2*X, X > 3) == \ + sqrt(170)*Rational(9, 85) + +def test_exponential_power(): + mu = Symbol('mu') + z = Symbol('z') + alpha = Symbol('alpha', positive=True) + beta = Symbol('beta', positive=True) + + X = ExponentialPower('x', mu, alpha, beta) + + assert density(X)(z) == beta*exp(-(Abs(mu - z)/alpha) + ** beta)/(2*alpha*gamma(1/beta)) + assert cdf(X)(z) == S.Half + lowergamma(1/beta, + (Abs(mu - z)/alpha)**beta)*sign(-mu + z)/\ + (2*gamma(1/beta)) + + +def test_f_distribution(): + d1 = Symbol("d1", positive=True) + d2 = Symbol("d2", positive=True) + + X = FDistribution("x", d1, d2) + + assert density(X)(x) == (d2**(d2/2)*sqrt((d1*x)**d1*(d1*x + d2)**(-d1 - d2)) + /(x*beta(d1/2, d2/2))) + + raises(NotImplementedError, lambda: moment_generating_function(X)) + d1 = Symbol("d1", nonpositive=True) + raises(ValueError, lambda: FDistribution('x', d1, d1)) + + d1 = Symbol("d1", positive=True, integer=False) + raises(ValueError, lambda: FDistribution('x', d1, d1)) + + d1 = Symbol("d1", positive=True) + d2 = Symbol("d2", nonpositive=True) + raises(ValueError, lambda: FDistribution('x', d1, d2)) + + d2 = Symbol("d2", positive=True, integer=False) + raises(ValueError, lambda: FDistribution('x', d1, d2)) + + +def test_fisher_z(): + d1 = Symbol("d1", positive=True) + d2 = Symbol("d2", positive=True) + + X = FisherZ("x", d1, d2) + assert density(X)(x) == (2*d1**(d1/2)*d2**(d2/2)*(d1*exp(2*x) + d2) + **(-d1/2 - d2/2)*exp(d1*x)/beta(d1/2, d2/2)) + +def test_frechet(): + a = Symbol("a", positive=True) + s = Symbol("s", positive=True) + m = Symbol("m", real=True) + + X = Frechet("x", a, s=s, m=m) + assert density(X)(x) == a*((x - m)/s)**(-a - 1)*exp(-((x - m)/s)**(-a))/s + assert cdf(X)(x) == Piecewise((exp(-((-m + x)/s)**(-a)), m <= x), (0, True)) + +@slow +def test_gamma(): + k = Symbol("k", positive=True) + theta = Symbol("theta", positive=True) + + X = Gamma('x', k, theta) + + # Tests characteristic function + assert characteristic_function(X)(x) == ((-I*theta*x + 1)**(-k)) + + assert density(X)(x) == x**(k - 1)*theta**(-k)*exp(-x/theta)/gamma(k) + assert cdf(X, meijerg=True)(z) == Piecewise( + (-k*lowergamma(k, 0)/gamma(k + 1) + + k*lowergamma(k, z/theta)/gamma(k + 1), z >= 0), + (0, True)) + + # assert simplify(variance(X)) == k*theta**2 # handled numerically below + assert E(X) == moment(X, 1) + + k, theta = symbols('k theta', positive=True) + X = Gamma('x', k, theta) + assert E(X) == k*theta + assert variance(X) == k*theta**2 + assert skewness(X).expand() == 2/sqrt(k) + assert kurtosis(X).expand() == 3 + 6/k + + Y = Gamma('y', 2*k, 3*theta) + assert coskewness(X, theta*X + Y, k*X + Y).simplify() == \ + 2*531441**(-k)*sqrt(k)*theta*(3*3**(12*k) - 2*531441**k) \ + /(sqrt(k**2 + 18)*sqrt(theta**2 + 18)) + +def test_gamma_inverse(): + a = Symbol("a", positive=True) + b = Symbol("b", positive=True) + X = GammaInverse("x", a, b) + assert density(X)(x) == x**(-a - 1)*b**a*exp(-b/x)/gamma(a) + assert cdf(X)(x) == Piecewise((uppergamma(a, b/x)/gamma(a), x > 0), (0, True)) + assert characteristic_function(X)(x) == 2 * (-I*b*x)**(a/2) \ + * besselk(a, 2*sqrt(b)*sqrt(-I*x))/gamma(a) + raises(NotImplementedError, lambda: moment_generating_function(X)) + +def test_gompertz(): + b = Symbol("b", positive=True) + eta = Symbol("eta", positive=True) + + X = Gompertz("x", b, eta) + + assert density(X)(x) == b*eta*exp(eta)*exp(b*x)*exp(-eta*exp(b*x)) + assert cdf(X)(x) == 1 - exp(eta)*exp(-eta*exp(b*x)) + assert diff(cdf(X)(x), x) == density(X)(x) + + +def test_gumbel(): + beta = Symbol("beta", positive=True) + mu = Symbol("mu") + x = Symbol("x") + y = Symbol("y") + X = Gumbel("x", beta, mu) + Y = Gumbel("y", beta, mu, minimum=True) + assert density(X)(x).expand() == \ + exp(mu/beta)*exp(-x/beta)*exp(-exp(mu/beta)*exp(-x/beta))/beta + assert density(Y)(y).expand() == \ + exp(-mu/beta)*exp(y/beta)*exp(-exp(-mu/beta)*exp(y/beta))/beta + assert cdf(X)(x).expand() == \ + exp(-exp(mu/beta)*exp(-x/beta)) + assert characteristic_function(X)(x) == exp(I*mu*x)*gamma(-I*beta*x + 1) + +def test_kumaraswamy(): + a = Symbol("a", positive=True) + b = Symbol("b", positive=True) + + X = Kumaraswamy("x", a, b) + assert density(X)(x) == x**(a - 1)*a*b*(-x**a + 1)**(b - 1) + assert cdf(X)(x) == Piecewise((0, x < 0), + (-(-x**a + 1)**b + 1, x <= 1), + (1, True)) + + +def test_laplace(): + mu = Symbol("mu") + b = Symbol("b", positive=True) + + X = Laplace('x', mu, b) + + #Tests characteristic_function + assert characteristic_function(X)(x) == (exp(I*mu*x)/(b**2*x**2 + 1)) + + assert density(X)(x) == exp(-Abs(x - mu)/b)/(2*b) + assert cdf(X)(x) == Piecewise((exp((-mu + x)/b)/2, mu > x), + (-exp((mu - x)/b)/2 + 1, True)) + X = Laplace('x', [1, 2], [[1, 0], [0, 1]]) + assert isinstance(pspace(X).distribution, MultivariateLaplaceDistribution) + +def test_levy(): + mu = Symbol("mu", real=True) + c = Symbol("c", positive=True) + + X = Levy('x', mu, c) + assert X.pspace.domain.set == Interval(mu, oo) + assert density(X)(x) == sqrt(c/(2*pi))*exp(-c/(2*(x - mu)))/((x - mu)**(S.One + S.Half)) + assert cdf(X)(x) == erfc(sqrt(c/(2*(x - mu)))) + + raises(NotImplementedError, lambda: moment_generating_function(X)) + mu = Symbol("mu", real=False) + raises(ValueError, lambda: Levy('x',mu,c)) + + c = Symbol("c", nonpositive=True) + raises(ValueError, lambda: Levy('x',mu,c)) + + mu = Symbol("mu", real=True) + raises(ValueError, lambda: Levy('x',mu,c)) + +def test_logcauchy(): + mu = Symbol("mu", positive=True) + sigma = Symbol("sigma", positive=True) + + X = LogCauchy("x", mu, sigma) + + assert density(X)(x) == sigma/(x*pi*(sigma**2 + (-mu + log(x))**2)) + assert cdf(X)(x) == atan((log(x) - mu)/sigma)/pi + S.Half + + +def test_logistic(): + mu = Symbol("mu", real=True) + s = Symbol("s", positive=True) + p = Symbol("p", positive=True) + + X = Logistic('x', mu, s) + + #Tests characteristics_function + assert characteristic_function(X)(x) == \ + (Piecewise((pi*s*x*exp(I*mu*x)/sinh(pi*s*x), Ne(x, 0)), (1, True))) + + assert density(X)(x) == exp((-x + mu)/s)/(s*(exp((-x + mu)/s) + 1)**2) + assert cdf(X)(x) == 1/(exp((mu - x)/s) + 1) + assert quantile(X)(p) == mu - s*log(-S.One + 1/p) + +def test_loglogistic(): + a, b = symbols('a b') + assert LogLogistic('x', a, b) + + a = Symbol('a', negative=True) + b = Symbol('b', positive=True) + raises(ValueError, lambda: LogLogistic('x', a, b)) + + a = Symbol('a', positive=True) + b = Symbol('b', negative=True) + raises(ValueError, lambda: LogLogistic('x', a, b)) + + a, b, z, p = symbols('a b z p', positive=True) + X = LogLogistic('x', a, b) + assert density(X)(z) == b*(z/a)**(b - 1)/(a*((z/a)**b + 1)**2) + assert cdf(X)(z) == 1/(1 + (z/a)**(-b)) + assert quantile(X)(p) == a*(p/(1 - p))**(1/b) + + # Expectation + assert E(X) == Piecewise((S.NaN, b <= 1), (pi*a/(b*sin(pi/b)), True)) + b = symbols('b', prime=True) # b > 1 + X = LogLogistic('x', a, b) + assert E(X) == pi*a/(b*sin(pi/b)) + X = LogLogistic('x', 1, 2) + assert median(X) == FiniteSet(1) + +def test_logitnormal(): + mu = Symbol('mu', real=True) + s = Symbol('s', positive=True) + X = LogitNormal('x', mu, s) + x = Symbol('x') + + assert density(X)(x) == sqrt(2)*exp(-(-mu + log(x/(1 - x)))**2/(2*s**2))/(2*sqrt(pi)*s*x*(1 - x)) + assert cdf(X)(x) == erf(sqrt(2)*(-mu + log(x/(1 - x)))/(2*s))/2 + S(1)/2 + +def test_lognormal(): + mean = Symbol('mu', real=True) + std = Symbol('sigma', positive=True) + X = LogNormal('x', mean, std) + # The sympy integrator can't do this too well + #assert E(X) == exp(mean+std**2/2) + #assert variance(X) == (exp(std**2)-1) * exp(2*mean + std**2) + + # The sympy integrator can't do this too well + #assert E(X) == + raises(NotImplementedError, lambda: moment_generating_function(X)) + mu = Symbol("mu", real=True) + sigma = Symbol("sigma", positive=True) + + X = LogNormal('x', mu, sigma) + assert density(X)(x) == (sqrt(2)*exp(-(-mu + log(x))**2 + /(2*sigma**2))/(2*x*sqrt(pi)*sigma)) + # Tests cdf + assert cdf(X)(x) == Piecewise( + (erf(sqrt(2)*(-mu + log(x))/(2*sigma))/2 + + S(1)/2, x > 0), (0, True)) + + X = LogNormal('x', 0, 1) # Mean 0, standard deviation 1 + assert density(X)(x) == sqrt(2)*exp(-log(x)**2/2)/(2*x*sqrt(pi)) + + +def test_Lomax(): + a, l = symbols('a, l', negative=True) + raises(ValueError, lambda: Lomax('X', a, l)) + a, l = symbols('a, l', real=False) + raises(ValueError, lambda: Lomax('X', a, l)) + + a, l = symbols('a, l', positive=True) + X = Lomax('X', a, l) + assert X.pspace.domain.set == Interval(0, oo) + assert density(X)(x) == a*(1 + x/l)**(-a - 1)/l + assert cdf(X)(x) == Piecewise((1 - (1 + x/l)**(-a), x >= 0), (0, True)) + a = 3 + X = Lomax('X', a, l) + assert E(X) == l/2 + assert median(X) == FiniteSet(l*(-1 + 2**Rational(1, 3))) + assert variance(X) == 3*l**2/4 + + +def test_maxwell(): + a = Symbol("a", positive=True) + + X = Maxwell('x', a) + + assert density(X)(x) == (sqrt(2)*x**2*exp(-x**2/(2*a**2))/ + (sqrt(pi)*a**3)) + assert E(X) == 2*sqrt(2)*a/sqrt(pi) + assert variance(X) == -8*a**2/pi + 3*a**2 + assert cdf(X)(x) == erf(sqrt(2)*x/(2*a)) - sqrt(2)*x*exp(-x**2/(2*a**2))/(sqrt(pi)*a) + assert diff(cdf(X)(x), x) == density(X)(x) + + +@slow +def test_Moyal(): + mu = Symbol('mu',real=False) + sigma = Symbol('sigma', positive=True) + raises(ValueError, lambda: Moyal('M',mu, sigma)) + + mu = Symbol('mu', real=True) + sigma = Symbol('sigma', negative=True) + raises(ValueError, lambda: Moyal('M',mu, sigma)) + + sigma = Symbol('sigma', positive=True) + M = Moyal('M', mu, sigma) + assert density(M)(z) == sqrt(2)*exp(-exp((mu - z)/sigma)/2 + - (-mu + z)/(2*sigma))/(2*sqrt(pi)*sigma) + assert cdf(M)(z).simplify() == 1 - erf(sqrt(2)*exp((mu - z)/(2*sigma))/2) + assert characteristic_function(M)(z) == 2**(-I*sigma*z)*exp(I*mu*z) \ + *gamma(-I*sigma*z + Rational(1, 2))/sqrt(pi) + assert E(M) == mu + EulerGamma*sigma + sigma*log(2) + assert moment_generating_function(M)(z) == 2**(-sigma*z)*exp(mu*z) \ + *gamma(-sigma*z + Rational(1, 2))/sqrt(pi) + + +def test_nakagami(): + mu = Symbol("mu", positive=True) + omega = Symbol("omega", positive=True) + + X = Nakagami('x', mu, omega) + assert density(X)(x) == (2*x**(2*mu - 1)*mu**mu*omega**(-mu) + *exp(-x**2*mu/omega)/gamma(mu)) + assert simplify(E(X)) == (sqrt(mu)*sqrt(omega) + *gamma(mu + S.Half)/gamma(mu + 1)) + assert simplify(variance(X)) == ( + omega - omega*gamma(mu + S.Half)**2/(gamma(mu)*gamma(mu + 1))) + assert cdf(X)(x) == Piecewise( + (lowergamma(mu, mu*x**2/omega)/gamma(mu), x > 0), + (0, True)) + X = Nakagami('x', 1, 1) + assert median(X) == FiniteSet(sqrt(log(2))) + +def test_gaussian_inverse(): + # test for symbolic parameters + a, b = symbols('a b') + assert GaussianInverse('x', a, b) + + # Inverse Gaussian distribution is also known as Wald distribution + # `GaussianInverse` can also be referred by the name `Wald` + a, b, z = symbols('a b z') + X = Wald('x', a, b) + assert density(X)(z) == sqrt(2)*sqrt(b/z**3)*exp(-b*(-a + z)**2/(2*a**2*z))/(2*sqrt(pi)) + + a, b = symbols('a b', positive=True) + z = Symbol('z', positive=True) + + X = GaussianInverse('x', a, b) + assert density(X)(z) == sqrt(2)*sqrt(b)*sqrt(z**(-3))*exp(-b*(-a + z)**2/(2*a**2*z))/(2*sqrt(pi)) + assert E(X) == a + assert variance(X).expand() == a**3/b + assert cdf(X)(z) == (S.Half - erf(sqrt(2)*sqrt(b)*(1 + z/a)/(2*sqrt(z)))/2)*exp(2*b/a) +\ + erf(sqrt(2)*sqrt(b)*(-1 + z/a)/(2*sqrt(z)))/2 + S.Half + + a = symbols('a', nonpositive=True) + raises(ValueError, lambda: GaussianInverse('x', a, b)) + + a = symbols('a', positive=True) + b = symbols('b', nonpositive=True) + raises(ValueError, lambda: GaussianInverse('x', a, b)) + +def test_pareto(): + xm, beta = symbols('xm beta', positive=True) + alpha = beta + 5 + X = Pareto('x', xm, alpha) + + dens = density(X) + + #Tests cdf function + assert cdf(X)(x) == \ + Piecewise((-x**(-beta - 5)*xm**(beta + 5) + 1, x >= xm), (0, True)) + + #Tests characteristic_function + assert characteristic_function(X)(x) == \ + ((-I*x*xm)**(beta + 5)*(beta + 5)*uppergamma(-beta - 5, -I*x*xm)) + + assert dens(x) == x**(-(alpha + 1))*xm**(alpha)*(alpha) + + assert simplify(E(X)) == alpha*xm/(alpha-1) + + # computation of taylor series for MGF still too slow + #assert simplify(variance(X)) == xm**2*alpha / ((alpha-1)**2*(alpha-2)) + + +def test_pareto_numeric(): + xm, beta = 3, 2 + alpha = beta + 5 + X = Pareto('x', xm, alpha) + + assert E(X) == alpha*xm/S(alpha - 1) + assert variance(X) == xm**2*alpha / S((alpha - 1)**2*(alpha - 2)) + assert median(X) == FiniteSet(3*2**Rational(1, 7)) + # Skewness tests too slow. Try shortcutting function? + + +def test_PowerFunction(): + alpha = Symbol("alpha", nonpositive=True) + a, b = symbols('a, b', real=True) + raises (ValueError, lambda: PowerFunction('x', alpha, a, b)) + + a, b = symbols('a, b', real=False) + raises (ValueError, lambda: PowerFunction('x', alpha, a, b)) + + alpha = Symbol("alpha", positive=True) + a, b = symbols('a, b', real=True) + raises (ValueError, lambda: PowerFunction('x', alpha, 5, 2)) + + X = PowerFunction('X', 2, a, b) + assert density(X)(z) == (-2*a + 2*z)/(-a + b)**2 + assert cdf(X)(z) == Piecewise((a**2/(a**2 - 2*a*b + b**2) - + 2*a*z/(a**2 - 2*a*b + b**2) + z**2/(a**2 - 2*a*b + b**2), a <= z), (0, True)) + + X = PowerFunction('X', 2, 0, 1) + assert density(X)(z) == 2*z + assert cdf(X)(z) == Piecewise((z**2, z >= 0), (0,True)) + assert E(X) == Rational(2,3) + assert P(X < 0) == 0 + assert P(X < 1) == 1 + assert median(X) == FiniteSet(1/sqrt(2)) + +def test_raised_cosine(): + mu = Symbol("mu", real=True) + s = Symbol("s", positive=True) + + X = RaisedCosine("x", mu, s) + + assert pspace(X).domain.set == Interval(mu - s, mu + s) + #Tests characteristics_function + assert characteristic_function(X)(x) == \ + Piecewise((exp(-I*pi*mu/s)/2, Eq(x, -pi/s)), (exp(I*pi*mu/s)/2, Eq(x, pi/s)), (pi**2*exp(I*mu*x)*sin(s*x)/(s*x*(-s**2*x**2 + pi**2)), True)) + + assert density(X)(x) == (Piecewise(((cos(pi*(x - mu)/s) + 1)/(2*s), + And(x <= mu + s, mu - s <= x)), (0, True))) + + +def test_rayleigh(): + sigma = Symbol("sigma", positive=True) + + X = Rayleigh('x', sigma) + + #Tests characteristic_function + assert characteristic_function(X)(x) == (-sqrt(2)*sqrt(pi)*sigma*x*(erfi(sqrt(2)*sigma*x/2) - I)*exp(-sigma**2*x**2/2)/2 + 1) + + assert density(X)(x) == x*exp(-x**2/(2*sigma**2))/sigma**2 + assert E(X) == sqrt(2)*sqrt(pi)*sigma/2 + assert variance(X) == -pi*sigma**2/2 + 2*sigma**2 + assert cdf(X)(x) == 1 - exp(-x**2/(2*sigma**2)) + assert diff(cdf(X)(x), x) == density(X)(x) + +def test_reciprocal(): + a = Symbol("a", real=True) + b = Symbol("b", real=True) + + X = Reciprocal('x', a, b) + assert density(X)(x) == 1/(x*(-log(a) + log(b))) + assert cdf(X)(x) == Piecewise((log(a)/(log(a) - log(b)) - log(x)/(log(a) - log(b)), a <= x), (0, True)) + X = Reciprocal('x', 5, 30) + + assert E(X) == 25/(log(30) - log(5)) + assert P(X < 4) == S.Zero + assert P(X < 20) == log(20) / (log(30) - log(5)) - log(5) / (log(30) - log(5)) + assert cdf(X)(10) == log(10) / (log(30) - log(5)) - log(5) / (log(30) - log(5)) + + a = symbols('a', nonpositive=True) + raises(ValueError, lambda: Reciprocal('x', a, b)) + + a = symbols('a', positive=True) + b = symbols('b', positive=True) + raises(ValueError, lambda: Reciprocal('x', a + b, a)) + +def test_shiftedgompertz(): + b = Symbol("b", positive=True) + eta = Symbol("eta", positive=True) + X = ShiftedGompertz("x", b, eta) + assert density(X)(x) == b*(eta*(1 - exp(-b*x)) + 1)*exp(-b*x)*exp(-eta*exp(-b*x)) + + +def test_studentt(): + nu = Symbol("nu", positive=True) + + X = StudentT('x', nu) + assert density(X)(x) == (1 + x**2/nu)**(-nu/2 - S.Half)/(sqrt(nu)*beta(S.Half, nu/2)) + assert cdf(X)(x) == S.Half + x*gamma(nu/2 + S.Half)*hyper((S.Half, nu/2 + S.Half), + (Rational(3, 2),), -x**2/nu)/(sqrt(pi)*sqrt(nu)*gamma(nu/2)) + raises(NotImplementedError, lambda: moment_generating_function(X)) + +def test_trapezoidal(): + a = Symbol("a", real=True) + b = Symbol("b", real=True) + c = Symbol("c", real=True) + d = Symbol("d", real=True) + + X = Trapezoidal('x', a, b, c, d) + assert density(X)(x) == Piecewise(((-2*a + 2*x)/((-a + b)*(-a - b + c + d)), (a <= x) & (x < b)), + (2/(-a - b + c + d), (b <= x) & (x < c)), + ((2*d - 2*x)/((-c + d)*(-a - b + c + d)), (c <= x) & (x <= d)), + (0, True)) + + X = Trapezoidal('x', 0, 1, 2, 3) + assert E(X) == Rational(3, 2) + assert variance(X) == Rational(5, 12) + assert P(X < 2) == Rational(3, 4) + assert median(X) == FiniteSet(Rational(3, 2)) + +def test_triangular(): + a = Symbol("a") + b = Symbol("b") + c = Symbol("c") + + X = Triangular('x', a, b, c) + assert pspace(X).domain.set == Interval(a, b) + assert str(density(X)(x)) == ("Piecewise(((-2*a + 2*x)/((-a + b)*(-a + c)), (a <= x) & (c > x)), " + "(2/(-a + b), Eq(c, x)), ((2*b - 2*x)/((-a + b)*(b - c)), (b >= x) & (c < x)), (0, True))") + + #Tests moment_generating_function + assert moment_generating_function(X)(x).expand() == \ + ((-2*(-a + b)*exp(c*x) + 2*(-a + c)*exp(b*x) + 2*(b - c)*exp(a*x))/(x**2*(-a + b)*(-a + c)*(b - c))).expand() + assert str(characteristic_function(X)(x)) == \ + '(2*(-a + b)*exp(I*c*x) - 2*(-a + c)*exp(I*b*x) - 2*(b - c)*exp(I*a*x))/(x**2*(-a + b)*(-a + c)*(b - c))' + +def test_quadratic_u(): + a = Symbol("a", real=True) + b = Symbol("b", real=True) + + X = QuadraticU("x", a, b) + Y = QuadraticU("x", 1, 2) + + assert pspace(X).domain.set == Interval(a, b) + # Tests _moment_generating_function + assert moment_generating_function(Y)(1) == -15*exp(2) + 27*exp(1) + assert moment_generating_function(Y)(2) == -9*exp(4)/2 + 21*exp(2)/2 + + assert characteristic_function(Y)(1) == 3*I*(-1 + 4*I)*exp(I*exp(2*I)) + assert density(X)(x) == (Piecewise((12*(x - a/2 - b/2)**2/(-a + b)**3, + And(x <= b, a <= x)), (0, True))) + + +def test_uniform(): + l = Symbol('l', real=True) + w = Symbol('w', positive=True) + X = Uniform('x', l, l + w) + + assert E(X) == l + w/2 + assert variance(X).expand() == w**2/12 + + # With numbers all is well + X = Uniform('x', 3, 5) + assert P(X < 3) == 0 and P(X > 5) == 0 + assert P(X < 4) == P(X > 4) == S.Half + assert median(X) == FiniteSet(4) + + z = Symbol('z') + p = density(X)(z) + assert p.subs(z, 3.7) == S.Half + assert p.subs(z, -1) == 0 + assert p.subs(z, 6) == 0 + + c = cdf(X) + assert c(2) == 0 and c(3) == 0 + assert c(Rational(7, 2)) == Rational(1, 4) + assert c(5) == 1 and c(6) == 1 + + +@XFAIL +@slow +def test_uniform_P(): + """ This stopped working because SingleContinuousPSpace.compute_density no + longer calls integrate on a DiracDelta but rather just solves directly. + integrate used to call UniformDistribution.expectation which special-cased + subsed out the Min and Max terms that Uniform produces + + I decided to regress on this class for general cleanliness (and I suspect + speed) of the algorithm. + """ + l = Symbol('l', real=True) + w = Symbol('w', positive=True) + X = Uniform('x', l, l + w) + assert P(X < l) == 0 and P(X > l + w) == 0 + + +def test_uniformsum(): + n = Symbol("n", integer=True) + _k = Dummy("k") + x = Symbol("x") + + X = UniformSum('x', n) + res = Sum((-1)**_k*(-_k + x)**(n - 1)*binomial(n, _k), (_k, 0, floor(x)))/factorial(n - 1) + assert density(X)(x).dummy_eq(res) + + #Tests set functions + assert X.pspace.domain.set == Interval(0, n) + + #Tests the characteristic_function + assert characteristic_function(X)(x) == (-I*(exp(I*x) - 1)/x)**n + + #Tests the moment_generating_function + assert moment_generating_function(X)(x) == ((exp(x) - 1)/x)**n + + +def test_von_mises(): + mu = Symbol("mu") + k = Symbol("k", positive=True) + + X = VonMises("x", mu, k) + assert density(X)(x) == exp(k*cos(x - mu))/(2*pi*besseli(0, k)) + + +def test_weibull(): + a, b = symbols('a b', positive=True) + # FIXME: simplify(E(X)) seems to hang without extended_positive=True + # On a Linux machine this had a rapid memory leak... + # a, b = symbols('a b', positive=True) + X = Weibull('x', a, b) + + assert E(X).expand() == a * gamma(1 + 1/b) + assert variance(X).expand() == (a**2 * gamma(1 + 2/b) - E(X)**2).expand() + assert simplify(skewness(X)) == (2*gamma(1 + 1/b)**3 - 3*gamma(1 + 1/b)*gamma(1 + 2/b) + gamma(1 + 3/b))/(-gamma(1 + 1/b)**2 + gamma(1 + 2/b))**Rational(3, 2) + assert simplify(kurtosis(X)) == (-3*gamma(1 + 1/b)**4 +\ + 6*gamma(1 + 1/b)**2*gamma(1 + 2/b) - 4*gamma(1 + 1/b)*gamma(1 + 3/b) + gamma(1 + 4/b))/(gamma(1 + 1/b)**2 - gamma(1 + 2/b))**2 + +def test_weibull_numeric(): + # Test for integers and rationals + a = 1 + bvals = [S.Half, 1, Rational(3, 2), 5] + for b in bvals: + X = Weibull('x', a, b) + assert simplify(E(X)) == expand_func(a * gamma(1 + 1/S(b))) + assert simplify(variance(X)) == simplify( + a**2 * gamma(1 + 2/S(b)) - E(X)**2) + # Not testing Skew... it's slow with int/frac values > 3/2 + + +def test_wignersemicircle(): + R = Symbol("R", positive=True) + + X = WignerSemicircle('x', R) + assert pspace(X).domain.set == Interval(-R, R) + assert density(X)(x) == 2*sqrt(-x**2 + R**2)/(pi*R**2) + assert E(X) == 0 + + + #Tests ChiNoncentralDistribution + assert characteristic_function(X)(x) == \ + Piecewise((2*besselj(1, R*x)/(R*x), Ne(x, 0)), (1, True)) + + +def test_input_value_assertions(): + a, b = symbols('a b') + p, q = symbols('p q', positive=True) + m, n = symbols('m n', positive=False, real=True) + + raises(ValueError, lambda: Normal('x', 3, 0)) + raises(ValueError, lambda: Normal('x', m, n)) + Normal('X', a, p) # No error raised + raises(ValueError, lambda: Exponential('x', m)) + Exponential('Ex', p) # No error raised + for fn in [Pareto, Weibull, Beta, Gamma]: + raises(ValueError, lambda: fn('x', m, p)) + raises(ValueError, lambda: fn('x', p, n)) + fn('x', p, q) # No error raised + + +def test_unevaluated(): + X = Normal('x', 0, 1) + k = Dummy('k') + expr1 = Integral(sqrt(2)*k*exp(-k**2/2)/(2*sqrt(pi)), (k, -oo, oo)) + expr2 = Integral(sqrt(2)*exp(-k**2/2)/(2*sqrt(pi)), (k, 0, oo)) + with ignore_warnings(UserWarning): ### TODO: Restore tests once warnings are removed + assert E(X, evaluate=False).rewrite(Integral).dummy_eq(expr1) + assert E(X + 1, evaluate=False).rewrite(Integral).dummy_eq(expr1 + 1) + assert P(X > 0, evaluate=False).rewrite(Integral).dummy_eq(expr2) + + assert P(X > 0, X**2 < 1) == S.Half + + +def test_probability_unevaluated(): + T = Normal('T', 30, 3) + with ignore_warnings(UserWarning): ### TODO: Restore tests once warnings are removed + assert type(P(T > 33, evaluate=False)) == Probability + + +def test_density_unevaluated(): + X = Normal('X', 0, 1) + Y = Normal('Y', 0, 2) + assert isinstance(density(X+Y, evaluate=False)(z), Integral) + + +def test_NormalDistribution(): + nd = NormalDistribution(0, 1) + x = Symbol('x') + assert nd.cdf(x) == erf(sqrt(2)*x/2)/2 + S.Half + assert nd.expectation(1, x) == 1 + assert nd.expectation(x, x) == 0 + assert nd.expectation(x**2, x) == 1 + #Test issue 10076 + a = SingleContinuousPSpace(x, NormalDistribution(2, 4)) + _z = Dummy('_z') + + expected1 = Integral(sqrt(2)*exp(-(_z - 2)**2/32)/(8*sqrt(pi)),(_z, -oo, 1)) + assert a.probability(x < 1, evaluate=False).dummy_eq(expected1) is True + + expected2 = Integral(sqrt(2)*exp(-(_z - 2)**2/32)/(8*sqrt(pi)),(_z, 1, oo)) + assert a.probability(x > 1, evaluate=False).dummy_eq(expected2) is True + + b = SingleContinuousPSpace(x, NormalDistribution(1, 9)) + + expected3 = Integral(sqrt(2)*exp(-(_z - 1)**2/162)/(18*sqrt(pi)),(_z, 6, oo)) + assert b.probability(x > 6, evaluate=False).dummy_eq(expected3) is True + + expected4 = Integral(sqrt(2)*exp(-(_z - 1)**2/162)/(18*sqrt(pi)),(_z, -oo, 6)) + assert b.probability(x < 6, evaluate=False).dummy_eq(expected4) is True + + +def test_random_parameters(): + mu = Normal('mu', 2, 3) + meas = Normal('T', mu, 1) + assert density(meas, evaluate=False)(z) + assert isinstance(pspace(meas), CompoundPSpace) + X = Normal('x', [1, 2], [[1, 0], [0, 1]]) + assert isinstance(pspace(X).distribution, MultivariateNormalDistribution) + assert density(meas)(z).simplify() == sqrt(5)*exp(-z**2/20 + z/5 - S(1)/5)/(10*sqrt(pi)) + + +def test_random_parameters_given(): + mu = Normal('mu', 2, 3) + meas = Normal('T', mu, 1) + assert given(meas, Eq(mu, 5)) == Normal('T', 5, 1) + + +def test_conjugate_priors(): + mu = Normal('mu', 2, 3) + x = Normal('x', mu, 1) + assert isinstance(simplify(density(mu, Eq(x, y), evaluate=False)(z)), + Mul) + + +def test_difficult_univariate(): + """ Since using solve in place of deltaintegrate we're able to perform + substantially more complex density computations on single continuous random + variables """ + x = Normal('x', 0, 1) + assert density(x**3) + assert density(exp(x**2)) + assert density(log(x)) + + +def test_issue_10003(): + X = Exponential('x', 3) + G = Gamma('g', 1, 2) + assert P(X < -1) is S.Zero + assert P(G < -1) is S.Zero + + +def test_precomputed_cdf(): + x = symbols("x", real=True) + mu = symbols("mu", real=True) + sigma, xm, alpha = symbols("sigma xm alpha", positive=True) + n = symbols("n", integer=True, positive=True) + distribs = [ + Normal("X", mu, sigma), + Pareto("P", xm, alpha), + ChiSquared("C", n), + Exponential("E", sigma), + # LogNormal("L", mu, sigma), + ] + for X in distribs: + compdiff = cdf(X)(x) - simplify(X.pspace.density.compute_cdf()(x)) + compdiff = simplify(compdiff.rewrite(erfc)) + assert compdiff == 0 + + +@slow +def test_precomputed_characteristic_functions(): + import mpmath + + def test_cf(dist, support_lower_limit, support_upper_limit): + pdf = density(dist) + t = Symbol('t') + + # first function is the hardcoded CF of the distribution + cf1 = lambdify([t], characteristic_function(dist)(t), 'mpmath') + + # second function is the Fourier transform of the density function + f = lambdify([x, t], pdf(x)*exp(I*x*t), 'mpmath') + cf2 = lambda t: mpmath.quad(lambda x: f(x, t), [support_lower_limit, support_upper_limit], maxdegree=10) + + # compare the two functions at various points + for test_point in [2, 5, 8, 11]: + n1 = cf1(test_point) + n2 = cf2(test_point) + + assert abs(re(n1) - re(n2)) < 1e-12 + assert abs(im(n1) - im(n2)) < 1e-12 + + test_cf(Beta('b', 1, 2), 0, 1) + test_cf(Chi('c', 3), 0, mpmath.inf) + test_cf(ChiSquared('c', 2), 0, mpmath.inf) + test_cf(Exponential('e', 6), 0, mpmath.inf) + test_cf(Logistic('l', 1, 2), -mpmath.inf, mpmath.inf) + test_cf(Normal('n', -1, 5), -mpmath.inf, mpmath.inf) + test_cf(RaisedCosine('r', 3, 1), 2, 4) + test_cf(Rayleigh('r', 0.5), 0, mpmath.inf) + test_cf(Uniform('u', -1, 1), -1, 1) + test_cf(WignerSemicircle('w', 3), -3, 3) + + +def test_long_precomputed_cdf(): + x = symbols("x", real=True) + distribs = [ + Arcsin("A", -5, 9), + Dagum("D", 4, 10, 3), + Erlang("E", 14, 5), + Frechet("F", 2, 6, -3), + Gamma("G", 2, 7), + GammaInverse("GI", 3, 5), + Kumaraswamy("K", 6, 8), + Laplace("LA", -5, 4), + Logistic("L", -6, 7), + Nakagami("N", 2, 7), + StudentT("S", 4) + ] + for distr in distribs: + for _ in range(5): + assert tn(diff(cdf(distr)(x), x), density(distr)(x), x, a=0, b=0, c=1, d=0) + + US = UniformSum("US", 5) + pdf01 = density(US)(x).subs(floor(x), 0).doit() # pdf on (0, 1) + cdf01 = cdf(US, evaluate=False)(x).subs(floor(x), 0).doit() # cdf on (0, 1) + assert tn(diff(cdf01, x), pdf01, x, a=0, b=0, c=1, d=0) + + +def test_issue_13324(): + X = Uniform('X', 0, 1) + assert E(X, X > S.Half) == Rational(3, 4) + assert E(X, X > 0) == S.Half + +def test_issue_20756(): + X = Uniform('X', -1, +1) + Y = Uniform('Y', -1, +1) + assert E(X * Y) == S.Zero + assert E(X * ((Y + 1) - 1)) == S.Zero + assert E(Y * (X*(X + 1) - X*X)) == S.Zero + +def test_FiniteSet_prob(): + E = Exponential('E', 3) + N = Normal('N', 5, 7) + assert P(Eq(E, 1)) is S.Zero + assert P(Eq(N, 2)) is S.Zero + assert P(Eq(N, x)) is S.Zero + +def test_prob_neq(): + E = Exponential('E', 4) + X = ChiSquared('X', 4) + assert P(Ne(E, 2)) == 1 + assert P(Ne(X, 4)) == 1 + assert P(Ne(X, 4)) == 1 + assert P(Ne(X, 5)) == 1 + assert P(Ne(E, x)) == 1 + +def test_union(): + N = Normal('N', 3, 2) + assert simplify(P(N**2 - N > 2)) == \ + -erf(sqrt(2))/2 - erfc(sqrt(2)/4)/2 + Rational(3, 2) + assert simplify(P(N**2 - 4 > 0)) == \ + -erf(5*sqrt(2)/4)/2 - erfc(sqrt(2)/4)/2 + Rational(3, 2) + +def test_Or(): + N = Normal('N', 0, 1) + assert simplify(P(Or(N > 2, N < 1))) == \ + -erf(sqrt(2))/2 - erfc(sqrt(2)/2)/2 + Rational(3, 2) + assert P(Or(N < 0, N < 1)) == P(N < 1) + assert P(Or(N > 0, N < 0)) == 1 + + +def test_conditional_eq(): + E = Exponential('E', 1) + assert P(Eq(E, 1), Eq(E, 1)) == 1 + assert P(Eq(E, 1), Eq(E, 2)) == 0 + assert P(E > 1, Eq(E, 2)) == 1 + assert P(E < 1, Eq(E, 2)) == 0 + +def test_ContinuousDistributionHandmade(): + x = Symbol('x') + z = Dummy('z') + dens = Lambda(x, Piecewise((S.Half, (0<=x)&(x<1)), (0, (x>=1)&(x<2)), + (S.Half, (x>=2)&(x<3)), (0, True))) + dens = ContinuousDistributionHandmade(dens, set=Interval(0, 3)) + space = SingleContinuousPSpace(z, dens) + assert dens.pdf == Lambda(x, Piecewise((S(1)/2, (x >= 0) & (x < 1)), + (0, (x >= 1) & (x < 2)), (S(1)/2, (x >= 2) & (x < 3)), (0, True))) + assert median(space.value) == Interval(1, 2) + assert E(space.value) == Rational(3, 2) + assert variance(space.value) == Rational(13, 12) + + +def test_issue_16318(): + # test compute_expectation function of the SingleContinuousDomain + N = SingleContinuousDomain(x, Interval(0, 1)) + raises(ValueError, lambda: SingleContinuousDomain.compute_expectation(N, x+1, {x, y})) + +def test_compute_density(): + X = Normal('X', 0, Symbol("sigma")**2) + raises(ValueError, lambda: density(X**5 + X)) diff --git a/lib/python3.12/site-packages/sympy/stats/tests/test_discrete_rv.py b/lib/python3.12/site-packages/sympy/stats/tests/test_discrete_rv.py new file mode 100644 index 0000000000000000000000000000000000000000..98303a460de2dac9e26278a7d06928ea616a9bce --- /dev/null +++ b/lib/python3.12/site-packages/sympy/stats/tests/test_discrete_rv.py @@ -0,0 +1,297 @@ +from sympy.concrete.summations import Sum +from sympy.core.numbers import (I, Rational, oo, pi) +from sympy.core.singleton import S +from sympy.core.symbol import Symbol +from sympy.functions.elementary.complexes import (im, re) +from sympy.functions.elementary.exponential import log +from sympy.functions.elementary.integers import floor +from sympy.functions.elementary.miscellaneous import sqrt +from sympy.functions.elementary.piecewise import Piecewise +from sympy.functions.special.bessel import besseli +from sympy.functions.special.beta_functions import beta +from sympy.functions.special.zeta_functions import zeta +from sympy.sets.sets import FiniteSet +from sympy.simplify.simplify import simplify +from sympy.utilities.lambdify import lambdify +from sympy.core.relational import Eq, Ne +from sympy.functions.elementary.exponential import exp +from sympy.logic.boolalg import Or +from sympy.sets.fancysets import Range +from sympy.stats import (P, E, variance, density, characteristic_function, + where, moment_generating_function, skewness, cdf, + kurtosis, coskewness) +from sympy.stats.drv_types import (PoissonDistribution, GeometricDistribution, + FlorySchulz, Poisson, Geometric, Hermite, Logarithmic, + NegativeBinomial, Skellam, YuleSimon, Zeta, + DiscreteRV) +from sympy.testing.pytest import slow, nocache_fail, raises +from sympy.stats.symbolic_probability import Expectation + +x = Symbol('x') + + +def test_PoissonDistribution(): + l = 3 + p = PoissonDistribution(l) + assert abs(p.cdf(10).evalf() - 1) < .001 + assert abs(p.cdf(10.4).evalf() - 1) < .001 + assert p.expectation(x, x) == l + assert p.expectation(x**2, x) - p.expectation(x, x)**2 == l + + +def test_Poisson(): + l = 3 + x = Poisson('x', l) + assert E(x) == l + assert E(2*x) == 2*l + assert variance(x) == l + assert density(x) == PoissonDistribution(l) + assert isinstance(E(x, evaluate=False), Expectation) + assert isinstance(E(2*x, evaluate=False), Expectation) + # issue 8248 + assert x.pspace.compute_expectation(1) == 1 + + +def test_FlorySchulz(): + a = Symbol("a") + z = Symbol("z") + x = FlorySchulz('x', a) + assert E(x) == (2 - a)/a + assert (variance(x) - 2*(1 - a)/a**2).simplify() == S(0) + assert density(x)(z) == a**2*z*(1 - a)**(z - 1) + + +@slow +def test_GeometricDistribution(): + p = S.One / 5 + d = GeometricDistribution(p) + assert d.expectation(x, x) == 1/p + assert d.expectation(x**2, x) - d.expectation(x, x)**2 == (1-p)/p**2 + assert abs(d.cdf(20000).evalf() - 1) < .001 + assert abs(d.cdf(20000.8).evalf() - 1) < .001 + G = Geometric('G', p=S(1)/4) + assert cdf(G)(S(7)/2) == P(G <= S(7)/2) + + X = Geometric('X', Rational(1, 5)) + Y = Geometric('Y', Rational(3, 10)) + assert coskewness(X, X + Y, X + 2*Y).simplify() == sqrt(230)*Rational(81, 1150) + + +def test_Hermite(): + a1 = Symbol("a1", positive=True) + a2 = Symbol("a2", negative=True) + raises(ValueError, lambda: Hermite("H", a1, a2)) + + a1 = Symbol("a1", negative=True) + a2 = Symbol("a2", positive=True) + raises(ValueError, lambda: Hermite("H", a1, a2)) + + a1 = Symbol("a1", positive=True) + x = Symbol("x") + H = Hermite("H", a1, a2) + assert moment_generating_function(H)(x) == exp(a1*(exp(x) - 1) + + a2*(exp(2*x) - 1)) + assert characteristic_function(H)(x) == exp(a1*(exp(I*x) - 1) + + a2*(exp(2*I*x) - 1)) + assert E(H) == a1 + 2*a2 + + H = Hermite("H", a1=5, a2=4) + assert density(H)(2) == 33*exp(-9)/2 + assert E(H) == 13 + assert variance(H) == 21 + assert kurtosis(H) == Rational(464,147) + assert skewness(H) == 37*sqrt(21)/441 + +def test_Logarithmic(): + p = S.Half + x = Logarithmic('x', p) + assert E(x) == -p / ((1 - p) * log(1 - p)) + assert variance(x) == -1/log(2)**2 + 2/log(2) + assert E(2*x**2 + 3*x + 4) == 4 + 7 / log(2) + assert isinstance(E(x, evaluate=False), Expectation) + + +@nocache_fail +def test_negative_binomial(): + r = 5 + p = S.One / 3 + x = NegativeBinomial('x', r, p) + assert E(x) == p*r / (1-p) + # This hangs when run with the cache disabled: + assert variance(x) == p*r / (1-p)**2 + assert E(x**5 + 2*x + 3) == Rational(9207, 4) + assert isinstance(E(x, evaluate=False), Expectation) + + +def test_skellam(): + mu1 = Symbol('mu1') + mu2 = Symbol('mu2') + z = Symbol('z') + X = Skellam('x', mu1, mu2) + + assert density(X)(z) == (mu1/mu2)**(z/2) * \ + exp(-mu1 - mu2)*besseli(z, 2*sqrt(mu1*mu2)) + assert skewness(X).expand() == mu1/(mu1*sqrt(mu1 + mu2) + mu2 * + sqrt(mu1 + mu2)) - mu2/(mu1*sqrt(mu1 + mu2) + mu2*sqrt(mu1 + mu2)) + assert variance(X).expand() == mu1 + mu2 + assert E(X) == mu1 - mu2 + assert characteristic_function(X)(z) == exp( + mu1*exp(I*z) - mu1 - mu2 + mu2*exp(-I*z)) + assert moment_generating_function(X)(z) == exp( + mu1*exp(z) - mu1 - mu2 + mu2*exp(-z)) + + +def test_yule_simon(): + from sympy.core.singleton import S + rho = S(3) + x = YuleSimon('x', rho) + assert simplify(E(x)) == rho / (rho - 1) + assert simplify(variance(x)) == rho**2 / ((rho - 1)**2 * (rho - 2)) + assert isinstance(E(x, evaluate=False), Expectation) + # To test the cdf function + assert cdf(x)(x) == Piecewise((-beta(floor(x), 4)*floor(x) + 1, x >= 1), (0, True)) + + +def test_zeta(): + s = S(5) + x = Zeta('x', s) + assert E(x) == zeta(s-1) / zeta(s) + assert simplify(variance(x)) == ( + zeta(s) * zeta(s-2) - zeta(s-1)**2) / zeta(s)**2 + + +def test_discrete_probability(): + X = Geometric('X', Rational(1, 5)) + Y = Poisson('Y', 4) + G = Geometric('e', x) + assert P(Eq(X, 3)) == Rational(16, 125) + assert P(X < 3) == Rational(9, 25) + assert P(X > 3) == Rational(64, 125) + assert P(X >= 3) == Rational(16, 25) + assert P(X <= 3) == Rational(61, 125) + assert P(Ne(X, 3)) == Rational(109, 125) + assert P(Eq(Y, 3)) == 32*exp(-4)/3 + assert P(Y < 3) == 13*exp(-4) + assert P(Y > 3).equals(32*(Rational(-71, 32) + 3*exp(4)/32)*exp(-4)/3) + assert P(Y >= 3).equals(32*(Rational(-39, 32) + 3*exp(4)/32)*exp(-4)/3) + assert P(Y <= 3) == 71*exp(-4)/3 + assert P(Ne(Y, 3)).equals( + 13*exp(-4) + 32*(Rational(-71, 32) + 3*exp(4)/32)*exp(-4)/3) + assert P(X < S.Infinity) is S.One + assert P(X > S.Infinity) is S.Zero + assert P(G < 3) == x*(2-x) + assert P(Eq(G, 3)) == x*(-x + 1)**2 + + +def test_DiscreteRV(): + p = S(1)/2 + x = Symbol('x', integer=True, positive=True) + pdf = p*(1 - p)**(x - 1) # pdf of Geometric Distribution + D = DiscreteRV(x, pdf, set=S.Naturals, check=True) + assert E(D) == E(Geometric('G', S(1)/2)) == 2 + assert P(D > 3) == S(1)/8 + assert D.pspace.domain.set == S.Naturals + raises(ValueError, lambda: DiscreteRV(x, x, FiniteSet(*range(4)), check=True)) + + # purposeful invalid pmf but it should not raise since check=False + # see test_drv_types.test_ContinuousRV for explanation + X = DiscreteRV(x, 1/x, S.Naturals) + assert P(X < 2) == 1 + assert E(X) == oo + +def test_precomputed_characteristic_functions(): + import mpmath + + def test_cf(dist, support_lower_limit, support_upper_limit): + pdf = density(dist) + t = S('t') + x = S('x') + + # first function is the hardcoded CF of the distribution + cf1 = lambdify([t], characteristic_function(dist)(t), 'mpmath') + + # second function is the Fourier transform of the density function + f = lambdify([x, t], pdf(x)*exp(I*x*t), 'mpmath') + cf2 = lambda t: mpmath.nsum(lambda x: f(x, t), [ + support_lower_limit, support_upper_limit], maxdegree=10) + + # compare the two functions at various points + for test_point in [2, 5, 8, 11]: + n1 = cf1(test_point) + n2 = cf2(test_point) + + assert abs(re(n1) - re(n2)) < 1e-12 + assert abs(im(n1) - im(n2)) < 1e-12 + + test_cf(Geometric('g', Rational(1, 3)), 1, mpmath.inf) + test_cf(Logarithmic('l', Rational(1, 5)), 1, mpmath.inf) + test_cf(NegativeBinomial('n', 5, Rational(1, 7)), 0, mpmath.inf) + test_cf(Poisson('p', 5), 0, mpmath.inf) + test_cf(YuleSimon('y', 5), 1, mpmath.inf) + test_cf(Zeta('z', 5), 1, mpmath.inf) + + +def test_moment_generating_functions(): + t = S('t') + + geometric_mgf = moment_generating_function(Geometric('g', S.Half))(t) + assert geometric_mgf.diff(t).subs(t, 0) == 2 + + logarithmic_mgf = moment_generating_function(Logarithmic('l', S.Half))(t) + assert logarithmic_mgf.diff(t).subs(t, 0) == 1/log(2) + + negative_binomial_mgf = moment_generating_function( + NegativeBinomial('n', 5, Rational(1, 3)))(t) + assert negative_binomial_mgf.diff(t).subs(t, 0) == Rational(5, 2) + + poisson_mgf = moment_generating_function(Poisson('p', 5))(t) + assert poisson_mgf.diff(t).subs(t, 0) == 5 + + skellam_mgf = moment_generating_function(Skellam('s', 1, 1))(t) + assert skellam_mgf.diff(t).subs( + t, 2) == (-exp(-2) + exp(2))*exp(-2 + exp(-2) + exp(2)) + + yule_simon_mgf = moment_generating_function(YuleSimon('y', 3))(t) + assert simplify(yule_simon_mgf.diff(t).subs(t, 0)) == Rational(3, 2) + + zeta_mgf = moment_generating_function(Zeta('z', 5))(t) + assert zeta_mgf.diff(t).subs(t, 0) == pi**4/(90*zeta(5)) + + +def test_Or(): + X = Geometric('X', S.Half) + assert P(Or(X < 3, X > 4)) == Rational(13, 16) + assert P(Or(X > 2, X > 1)) == P(X > 1) + assert P(Or(X >= 3, X < 3)) == 1 + + +def test_where(): + X = Geometric('X', Rational(1, 5)) + Y = Poisson('Y', 4) + assert where(X**2 > 4).set == Range(3, S.Infinity, 1) + assert where(X**2 >= 4).set == Range(2, S.Infinity, 1) + assert where(Y**2 < 9).set == Range(0, 3, 1) + assert where(Y**2 <= 9).set == Range(0, 4, 1) + + +def test_conditional(): + X = Geometric('X', Rational(2, 3)) + Y = Poisson('Y', 3) + assert P(X > 2, X > 3) == 1 + assert P(X > 3, X > 2) == Rational(1, 3) + assert P(Y > 2, Y < 2) == 0 + assert P(Eq(Y, 3), Y >= 0) == 9*exp(-3)/2 + assert P(Eq(Y, 3), Eq(Y, 2)) == 0 + assert P(X < 2, Eq(X, 2)) == 0 + assert P(X > 2, Eq(X, 3)) == 1 + + +def test_product_spaces(): + X1 = Geometric('X1', S.Half) + X2 = Geometric('X2', Rational(1, 3)) + assert str(P(X1 + X2 < 3).rewrite(Sum)) == ( + "Sum(Piecewise((1/(4*2**n), n >= -1), (0, True)), (n, -oo, -1))/3") + assert str(P(X1 + X2 > 3).rewrite(Sum)) == ( + 'Sum(Piecewise((2**(X2 - n - 2)*(3/2)**(1 - X2)/6, ' + 'X2 - n <= 2), (0, True)), (X2, 1, oo), (n, 1, oo))') + assert P(Eq(X1 + X2, 3)) == Rational(1, 12) diff --git a/lib/python3.12/site-packages/sympy/stats/tests/test_error_prop.py b/lib/python3.12/site-packages/sympy/stats/tests/test_error_prop.py new file mode 100644 index 0000000000000000000000000000000000000000..483fb4c36e202d744faeb355606ff9803a516873 --- /dev/null +++ b/lib/python3.12/site-packages/sympy/stats/tests/test_error_prop.py @@ -0,0 +1,60 @@ +from sympy.core.function import Function +from sympy.core.symbol import symbols +from sympy.functions.elementary.exponential import exp +from sympy.stats.error_prop import variance_prop +from sympy.stats.symbolic_probability import (RandomSymbol, Variance, + Covariance) + + +def test_variance_prop(): + x, y, z = symbols('x y z') + phi, t = consts = symbols('phi t') + a = RandomSymbol(x) + var_x = Variance(a) + var_y = Variance(RandomSymbol(y)) + var_z = Variance(RandomSymbol(z)) + f = Function('f')(x) + cases = { + x + y: var_x + var_y, + a + y: var_x + var_y, + x + y + z: var_x + var_y + var_z, + 2*x: 4*var_x, + x*y: var_x*y**2 + var_y*x**2, + 1/x: var_x/x**4, + x/y: (var_x*y**2 + var_y*x**2)/y**4, + exp(x): var_x*exp(2*x), + exp(2*x): 4*var_x*exp(4*x), + exp(-x*t): t**2*var_x*exp(-2*t*x), + f: Variance(f), + } + for inp, out in cases.items(): + obs = variance_prop(inp, consts=consts) + assert out == obs + +def test_variance_prop_with_covar(): + x, y, z = symbols('x y z') + phi, t = consts = symbols('phi t') + a = RandomSymbol(x) + var_x = Variance(a) + b = RandomSymbol(y) + var_y = Variance(b) + c = RandomSymbol(z) + var_z = Variance(c) + covar_x_y = Covariance(a, b) + covar_x_z = Covariance(a, c) + covar_y_z = Covariance(b, c) + cases = { + x + y: var_x + var_y + 2*covar_x_y, + a + y: var_x + var_y + 2*covar_x_y, + x + y + z: var_x + var_y + var_z + \ + 2*covar_x_y + 2*covar_x_z + 2*covar_y_z, + 2*x: 4*var_x, + x*y: var_x*y**2 + var_y*x**2 + 2*covar_x_y/(x*y), + 1/x: var_x/x**4, + exp(x): var_x*exp(2*x), + exp(2*x): 4*var_x*exp(4*x), + exp(-x*t): t**2*var_x*exp(-2*t*x), + } + for inp, out in cases.items(): + obs = variance_prop(inp, consts=consts, include_covar=True) + assert out == obs diff --git a/lib/python3.12/site-packages/sympy/stats/tests/test_finite_rv.py b/lib/python3.12/site-packages/sympy/stats/tests/test_finite_rv.py new file mode 100644 index 0000000000000000000000000000000000000000..93bf0211a26ecc32d7f18c7e2d8236859857e445 --- /dev/null +++ b/lib/python3.12/site-packages/sympy/stats/tests/test_finite_rv.py @@ -0,0 +1,509 @@ +from sympy.concrete.summations import Sum +from sympy.core.containers import (Dict, Tuple) +from sympy.core.function import Function +from sympy.core.numbers import (I, Rational, nan) +from sympy.core.relational import Eq +from sympy.core.singleton import S +from sympy.core.symbol import (Dummy, Symbol, symbols) +from sympy.core.sympify import sympify +from sympy.functions.combinatorial.factorials import binomial +from sympy.functions.combinatorial.numbers import harmonic +from sympy.functions.elementary.exponential import exp +from sympy.functions.elementary.miscellaneous import sqrt +from sympy.functions.elementary.piecewise import Piecewise +from sympy.functions.elementary.trigonometric import cos +from sympy.functions.special.beta_functions import beta +from sympy.logic.boolalg import (And, Or) +from sympy.polys.polytools import cancel +from sympy.sets.sets import FiniteSet +from sympy.simplify.simplify import simplify +from sympy.matrices import Matrix +from sympy.stats import (DiscreteUniform, Die, Bernoulli, Coin, Binomial, BetaBinomial, + Hypergeometric, Rademacher, IdealSoliton, RobustSoliton, P, E, variance, + covariance, skewness, density, where, FiniteRV, pspace, cdf, + correlation, moment, cmoment, smoment, characteristic_function, + moment_generating_function, quantile, kurtosis, median, coskewness) +from sympy.stats.frv_types import DieDistribution, BinomialDistribution, \ + HypergeometricDistribution +from sympy.stats.rv import Density +from sympy.testing.pytest import raises + + +def BayesTest(A, B): + assert P(A, B) == P(And(A, B)) / P(B) + assert P(A, B) == P(B, A) * P(A) / P(B) + + +def test_discreteuniform(): + # Symbolic + a, b, c, t = symbols('a b c t') + X = DiscreteUniform('X', [a, b, c]) + + assert E(X) == (a + b + c)/3 + assert simplify(variance(X) + - ((a**2 + b**2 + c**2)/3 - (a/3 + b/3 + c/3)**2)) == 0 + assert P(Eq(X, a)) == P(Eq(X, b)) == P(Eq(X, c)) == S('1/3') + + Y = DiscreteUniform('Y', range(-5, 5)) + + # Numeric + assert E(Y) == S('-1/2') + assert variance(Y) == S('33/4') + assert median(Y) == FiniteSet(-1, 0) + + for x in range(-5, 5): + assert P(Eq(Y, x)) == S('1/10') + assert P(Y <= x) == S(x + 6)/10 + assert P(Y >= x) == S(5 - x)/10 + + assert dict(density(Die('D', 6)).items()) == \ + dict(density(DiscreteUniform('U', range(1, 7))).items()) + + assert characteristic_function(X)(t) == exp(I*a*t)/3 + exp(I*b*t)/3 + exp(I*c*t)/3 + assert moment_generating_function(X)(t) == exp(a*t)/3 + exp(b*t)/3 + exp(c*t)/3 + # issue 18611 + raises(ValueError, lambda: DiscreteUniform('Z', [a, a, a, b, b, c])) + +def test_dice(): + # TODO: Make iid method! + X, Y, Z = Die('X', 6), Die('Y', 6), Die('Z', 6) + a, b, t, p = symbols('a b t p') + + assert E(X) == 3 + S.Half + assert variance(X) == Rational(35, 12) + assert E(X + Y) == 7 + assert E(X + X) == 7 + assert E(a*X + b) == a*E(X) + b + assert variance(X + Y) == variance(X) + variance(Y) == cmoment(X + Y, 2) + assert variance(X + X) == 4 * variance(X) == cmoment(X + X, 2) + assert cmoment(X, 0) == 1 + assert cmoment(4*X, 3) == 64*cmoment(X, 3) + assert covariance(X, Y) is S.Zero + assert covariance(X, X + Y) == variance(X) + assert density(Eq(cos(X*S.Pi), 1))[True] == S.Half + assert correlation(X, Y) == 0 + assert correlation(X, Y) == correlation(Y, X) + assert smoment(X + Y, 3) == skewness(X + Y) + assert smoment(X + Y, 4) == kurtosis(X + Y) + assert smoment(X, 0) == 1 + assert P(X > 3) == S.Half + assert P(2*X > 6) == S.Half + assert P(X > Y) == Rational(5, 12) + assert P(Eq(X, Y)) == P(Eq(X, 1)) + + assert E(X, X > 3) == 5 == moment(X, 1, 0, X > 3) + assert E(X, Y > 3) == E(X) == moment(X, 1, 0, Y > 3) + assert E(X + Y, Eq(X, Y)) == E(2*X) + assert moment(X, 0) == 1 + assert moment(5*X, 2) == 25*moment(X, 2) + assert quantile(X)(p) == Piecewise((nan, (p > 1) | (p < 0)),\ + (S.One, p <= Rational(1, 6)), (S(2), p <= Rational(1, 3)), (S(3), p <= S.Half),\ + (S(4), p <= Rational(2, 3)), (S(5), p <= Rational(5, 6)), (S(6), p <= 1)) + + assert P(X > 3, X > 3) is S.One + assert P(X > Y, Eq(Y, 6)) is S.Zero + assert P(Eq(X + Y, 12)) == Rational(1, 36) + assert P(Eq(X + Y, 12), Eq(X, 6)) == Rational(1, 6) + + assert density(X + Y) == density(Y + Z) != density(X + X) + d = density(2*X + Y**Z) + assert d[S(22)] == Rational(1, 108) and d[S(4100)] == Rational(1, 216) and S(3130) not in d + + assert pspace(X).domain.as_boolean() == Or( + *[Eq(X.symbol, i) for i in [1, 2, 3, 4, 5, 6]]) + + assert where(X > 3).set == FiniteSet(4, 5, 6) + + assert characteristic_function(X)(t) == exp(6*I*t)/6 + exp(5*I*t)/6 + exp(4*I*t)/6 + exp(3*I*t)/6 + exp(2*I*t)/6 + exp(I*t)/6 + assert moment_generating_function(X)(t) == exp(6*t)/6 + exp(5*t)/6 + exp(4*t)/6 + exp(3*t)/6 + exp(2*t)/6 + exp(t)/6 + assert median(X) == FiniteSet(3, 4) + D = Die('D', 7) + assert median(D) == FiniteSet(4) + # Bayes test for die + BayesTest(X > 3, X + Y < 5) + BayesTest(Eq(X - Y, Z), Z > Y) + BayesTest(X > 3, X > 2) + + # arg test for die + raises(ValueError, lambda: Die('X', -1)) # issue 8105: negative sides. + raises(ValueError, lambda: Die('X', 0)) + raises(ValueError, lambda: Die('X', 1.5)) # issue 8103: non integer sides. + + # symbolic test for die + n, k = symbols('n, k', positive=True) + D = Die('D', n) + dens = density(D).dict + assert dens == Density(DieDistribution(n)) + assert set(dens.subs(n, 4).doit().keys()) == {1, 2, 3, 4} + assert set(dens.subs(n, 4).doit().values()) == {Rational(1, 4)} + k = Dummy('k', integer=True) + assert E(D).dummy_eq( + Sum(Piecewise((k/n, k <= n), (0, True)), (k, 1, n))) + assert variance(D).subs(n, 6).doit() == Rational(35, 12) + + ki = Dummy('ki') + cumuf = cdf(D)(k) + assert cumuf.dummy_eq( + Sum(Piecewise((1/n, (ki >= 1) & (ki <= n)), (0, True)), (ki, 1, k))) + assert cumuf.subs({n: 6, k: 2}).doit() == Rational(1, 3) + + t = Dummy('t') + cf = characteristic_function(D)(t) + assert cf.dummy_eq( + Sum(Piecewise((exp(ki*I*t)/n, (ki >= 1) & (ki <= n)), (0, True)), (ki, 1, n))) + assert cf.subs(n, 3).doit() == exp(3*I*t)/3 + exp(2*I*t)/3 + exp(I*t)/3 + mgf = moment_generating_function(D)(t) + assert mgf.dummy_eq( + Sum(Piecewise((exp(ki*t)/n, (ki >= 1) & (ki <= n)), (0, True)), (ki, 1, n))) + assert mgf.subs(n, 3).doit() == exp(3*t)/3 + exp(2*t)/3 + exp(t)/3 + +def test_given(): + X = Die('X', 6) + assert density(X, X > 5) == {S(6): S.One} + assert where(X > 2, X > 5).as_boolean() == Eq(X.symbol, 6) + + +def test_domains(): + X, Y = Die('x', 6), Die('y', 6) + x, y = X.symbol, Y.symbol + # Domains + d = where(X > Y) + assert d.condition == (x > y) + d = where(And(X > Y, Y > 3)) + assert d.as_boolean() == Or(And(Eq(x, 5), Eq(y, 4)), And(Eq(x, 6), + Eq(y, 5)), And(Eq(x, 6), Eq(y, 4))) + assert len(d.elements) == 3 + + assert len(pspace(X + Y).domain.elements) == 36 + + Z = Die('x', 4) + + raises(ValueError, lambda: P(X > Z)) # Two domains with same internal symbol + + assert pspace(X + Y).domain.set == FiniteSet(1, 2, 3, 4, 5, 6)**2 + + assert where(X > 3).set == FiniteSet(4, 5, 6) + assert X.pspace.domain.dict == FiniteSet( + *[Dict({X.symbol: i}) for i in range(1, 7)]) + + assert where(X > Y).dict == FiniteSet(*[Dict({X.symbol: i, Y.symbol: j}) + for i in range(1, 7) for j in range(1, 7) if i > j]) + +def test_bernoulli(): + p, a, b, t = symbols('p a b t') + X = Bernoulli('B', p, a, b) + + assert E(X) == a*p + b*(-p + 1) + assert density(X)[a] == p + assert density(X)[b] == 1 - p + assert characteristic_function(X)(t) == p * exp(I * a * t) + (-p + 1) * exp(I * b * t) + assert moment_generating_function(X)(t) == p * exp(a * t) + (-p + 1) * exp(b * t) + + X = Bernoulli('B', p, 1, 0) + z = Symbol("z") + + assert E(X) == p + assert simplify(variance(X)) == p*(1 - p) + assert E(a*X + b) == a*E(X) + b + assert simplify(variance(a*X + b)) == simplify(a**2 * variance(X)) + assert quantile(X)(z) == Piecewise((nan, (z > 1) | (z < 0)), (0, z <= 1 - p), (1, z <= 1)) + Y = Bernoulli('Y', Rational(1, 2)) + assert median(Y) == FiniteSet(0, 1) + Z = Bernoulli('Z', Rational(2, 3)) + assert median(Z) == FiniteSet(1) + raises(ValueError, lambda: Bernoulli('B', 1.5)) + raises(ValueError, lambda: Bernoulli('B', -0.5)) + + #issue 8248 + assert X.pspace.compute_expectation(1) == 1 + + p = Rational(1, 5) + X = Binomial('X', 5, p) + Y = Binomial('Y', 7, 2*p) + Z = Binomial('Z', 9, 3*p) + assert coskewness(Y + Z, X + Y, X + Z).simplify() == 0 + assert coskewness(Y + 2*X + Z, X + 2*Y + Z, X + 2*Z + Y).simplify() == \ + sqrt(1529)*Rational(12, 16819) + assert coskewness(Y + 2*X + Z, X + 2*Y + Z, X + 2*Z + Y, X < 2).simplify() \ + == -sqrt(357451121)*Rational(2812, 4646864573) + +def test_cdf(): + D = Die('D', 6) + o = S.One + + assert cdf( + D) == sympify({1: o/6, 2: o/3, 3: o/2, 4: 2*o/3, 5: 5*o/6, 6: o}) + + +def test_coins(): + C, D = Coin('C'), Coin('D') + H, T = symbols('H, T') + assert P(Eq(C, D)) == S.Half + assert density(Tuple(C, D)) == {(H, H): Rational(1, 4), (H, T): Rational(1, 4), + (T, H): Rational(1, 4), (T, T): Rational(1, 4)} + assert dict(density(C).items()) == {H: S.Half, T: S.Half} + + F = Coin('F', Rational(1, 10)) + assert P(Eq(F, H)) == Rational(1, 10) + + d = pspace(C).domain + + assert d.as_boolean() == Or(Eq(C.symbol, H), Eq(C.symbol, T)) + + raises(ValueError, lambda: P(C > D)) # Can't intelligently compare H to T + +def test_binomial_verify_parameters(): + raises(ValueError, lambda: Binomial('b', .2, .5)) + raises(ValueError, lambda: Binomial('b', 3, 1.5)) + +def test_binomial_numeric(): + nvals = range(5) + pvals = [0, Rational(1, 4), S.Half, Rational(3, 4), 1] + + for n in nvals: + for p in pvals: + X = Binomial('X', n, p) + assert E(X) == n*p + assert variance(X) == n*p*(1 - p) + if n > 0 and 0 < p < 1: + assert skewness(X) == (1 - 2*p)/sqrt(n*p*(1 - p)) + assert kurtosis(X) == 3 + (1 - 6*p*(1 - p))/(n*p*(1 - p)) + for k in range(n + 1): + assert P(Eq(X, k)) == binomial(n, k)*p**k*(1 - p)**(n - k) + +def test_binomial_quantile(): + X = Binomial('X', 50, S.Half) + assert quantile(X)(0.95) == S(31) + assert median(X) == FiniteSet(25) + + X = Binomial('X', 5, S.Half) + p = Symbol("p", positive=True) + assert quantile(X)(p) == Piecewise((nan, p > S.One), (S.Zero, p <= Rational(1, 32)),\ + (S.One, p <= Rational(3, 16)), (S(2), p <= S.Half), (S(3), p <= Rational(13, 16)),\ + (S(4), p <= Rational(31, 32)), (S(5), p <= S.One)) + assert median(X) == FiniteSet(2, 3) + + +def test_binomial_symbolic(): + n = 2 + p = symbols('p', positive=True) + X = Binomial('X', n, p) + t = Symbol('t') + + assert simplify(E(X)) == n*p == simplify(moment(X, 1)) + assert simplify(variance(X)) == n*p*(1 - p) == simplify(cmoment(X, 2)) + assert cancel(skewness(X) - (1 - 2*p)/sqrt(n*p*(1 - p))) == 0 + assert cancel((kurtosis(X)) - (3 + (1 - 6*p*(1 - p))/(n*p*(1 - p)))) == 0 + assert characteristic_function(X)(t) == p ** 2 * exp(2 * I * t) + 2 * p * (-p + 1) * exp(I * t) + (-p + 1) ** 2 + assert moment_generating_function(X)(t) == p ** 2 * exp(2 * t) + 2 * p * (-p + 1) * exp(t) + (-p + 1) ** 2 + + # Test ability to change success/failure winnings + H, T = symbols('H T') + Y = Binomial('Y', n, p, succ=H, fail=T) + assert simplify(E(Y) - (n*(H*p + T*(1 - p)))) == 0 + + # test symbolic dimensions + n = symbols('n') + B = Binomial('B', n, p) + raises(NotImplementedError, lambda: P(B > 2)) + assert density(B).dict == Density(BinomialDistribution(n, p, 1, 0)) + assert set(density(B).dict.subs(n, 4).doit().keys()) == \ + {S.Zero, S.One, S(2), S(3), S(4)} + assert set(density(B).dict.subs(n, 4).doit().values()) == \ + {(1 - p)**4, 4*p*(1 - p)**3, 6*p**2*(1 - p)**2, 4*p**3*(1 - p), p**4} + k = Dummy('k', integer=True) + assert E(B > 2).dummy_eq( + Sum(Piecewise((k*p**k*(1 - p)**(-k + n)*binomial(n, k), (k >= 0) + & (k <= n) & (k > 2)), (0, True)), (k, 0, n))) + +def test_beta_binomial(): + # verify parameters + raises(ValueError, lambda: BetaBinomial('b', .2, 1, 2)) + raises(ValueError, lambda: BetaBinomial('b', 2, -1, 2)) + raises(ValueError, lambda: BetaBinomial('b', 2, 1, -2)) + assert BetaBinomial('b', 2, 1, 1) + + # test numeric values + nvals = range(1,5) + alphavals = [Rational(1, 4), S.Half, Rational(3, 4), 1, 10] + betavals = [Rational(1, 4), S.Half, Rational(3, 4), 1, 10] + + for n in nvals: + for a in alphavals: + for b in betavals: + X = BetaBinomial('X', n, a, b) + assert E(X) == moment(X, 1) + assert variance(X) == cmoment(X, 2) + + # test symbolic + n, a, b = symbols('a b n') + assert BetaBinomial('x', n, a, b) + n = 2 # Because we're using for loops, can't do symbolic n + a, b = symbols('a b', positive=True) + X = BetaBinomial('X', n, a, b) + t = Symbol('t') + + assert E(X).expand() == moment(X, 1).expand() + assert variance(X).expand() == cmoment(X, 2).expand() + assert skewness(X) == smoment(X, 3) + assert characteristic_function(X)(t) == exp(2*I*t)*beta(a + 2, b)/beta(a, b) +\ + 2*exp(I*t)*beta(a + 1, b + 1)/beta(a, b) + beta(a, b + 2)/beta(a, b) + assert moment_generating_function(X)(t) == exp(2*t)*beta(a + 2, b)/beta(a, b) +\ + 2*exp(t)*beta(a + 1, b + 1)/beta(a, b) + beta(a, b + 2)/beta(a, b) + +def test_hypergeometric_numeric(): + for N in range(1, 5): + for m in range(0, N + 1): + for n in range(1, N + 1): + X = Hypergeometric('X', N, m, n) + N, m, n = map(sympify, (N, m, n)) + assert sum(density(X).values()) == 1 + assert E(X) == n * m / N + if N > 1: + assert variance(X) == n*(m/N)*(N - m)/N*(N - n)/(N - 1) + # Only test for skewness when defined + if N > 2 and 0 < m < N and n < N: + assert skewness(X) == simplify((N - 2*m)*sqrt(N - 1)*(N - 2*n) + / (sqrt(n*m*(N - m)*(N - n))*(N - 2))) + +def test_hypergeometric_symbolic(): + N, m, n = symbols('N, m, n') + H = Hypergeometric('H', N, m, n) + dens = density(H).dict + expec = E(H > 2) + assert dens == Density(HypergeometricDistribution(N, m, n)) + assert dens.subs(N, 5).doit() == Density(HypergeometricDistribution(5, m, n)) + assert set(dens.subs({N: 3, m: 2, n: 1}).doit().keys()) == {S.Zero, S.One} + assert set(dens.subs({N: 3, m: 2, n: 1}).doit().values()) == {Rational(1, 3), Rational(2, 3)} + k = Dummy('k', integer=True) + assert expec.dummy_eq( + Sum(Piecewise((k*binomial(m, k)*binomial(N - m, -k + n) + /binomial(N, n), k > 2), (0, True)), (k, 0, n))) + +def test_rademacher(): + X = Rademacher('X') + t = Symbol('t') + + assert E(X) == 0 + assert variance(X) == 1 + assert density(X)[-1] == S.Half + assert density(X)[1] == S.Half + assert characteristic_function(X)(t) == exp(I*t)/2 + exp(-I*t)/2 + assert moment_generating_function(X)(t) == exp(t) / 2 + exp(-t) / 2 + +def test_ideal_soliton(): + raises(ValueError, lambda : IdealSoliton('sol', -12)) + raises(ValueError, lambda : IdealSoliton('sol', 13.2)) + raises(ValueError, lambda : IdealSoliton('sol', 0)) + f = Function('f') + raises(ValueError, lambda : density(IdealSoliton('sol', 10)).pmf(f)) + + k = Symbol('k', integer=True, positive=True) + x = Symbol('x', integer=True, positive=True) + t = Symbol('t') + sol = IdealSoliton('sol', k) + assert density(sol).low == S.One + assert density(sol).high == k + assert density(sol).dict == Density(density(sol)) + assert density(sol).pmf(x) == Piecewise((1/k, Eq(x, 1)), (1/(x*(x - 1)), k >= x), (0, True)) + + k_vals = [5, 20, 50, 100, 1000] + for i in k_vals: + assert E(sol.subs(k, i)) == harmonic(i) == moment(sol.subs(k, i), 1) + assert variance(sol.subs(k, i)) == (i - 1) + harmonic(i) - harmonic(i)**2 == cmoment(sol.subs(k, i),2) + assert skewness(sol.subs(k, i)) == smoment(sol.subs(k, i), 3) + assert kurtosis(sol.subs(k, i)) == smoment(sol.subs(k, i), 4) + + assert exp(I*t)/10 + Sum(exp(I*t*x)/(x*x - x), (x, 2, k)).subs(k, 10).doit() == characteristic_function(sol.subs(k, 10))(t) + assert exp(t)/10 + Sum(exp(t*x)/(x*x - x), (x, 2, k)).subs(k, 10).doit() == moment_generating_function(sol.subs(k, 10))(t) + +def test_robust_soliton(): + raises(ValueError, lambda : RobustSoliton('robSol', -12, 0.1, 0.02)) + raises(ValueError, lambda : RobustSoliton('robSol', 13, 1.89, 0.1)) + raises(ValueError, lambda : RobustSoliton('robSol', 15, 0.6, -2.31)) + f = Function('f') + raises(ValueError, lambda : density(RobustSoliton('robSol', 15, 0.6, 0.1)).pmf(f)) + + k = Symbol('k', integer=True, positive=True) + delta = Symbol('delta', positive=True) + c = Symbol('c', positive=True) + robSol = RobustSoliton('robSol', k, delta, c) + assert density(robSol).low == 1 + assert density(robSol).high == k + + k_vals = [10, 20, 50] + delta_vals = [0.2, 0.4, 0.6] + c_vals = [0.01, 0.03, 0.05] + for x in k_vals: + for y in delta_vals: + for z in c_vals: + assert E(robSol.subs({k: x, delta: y, c: z})) == moment(robSol.subs({k: x, delta: y, c: z}), 1) + assert variance(robSol.subs({k: x, delta: y, c: z})) == cmoment(robSol.subs({k: x, delta: y, c: z}), 2) + assert skewness(robSol.subs({k: x, delta: y, c: z})) == smoment(robSol.subs({k: x, delta: y, c: z}), 3) + assert kurtosis(robSol.subs({k: x, delta: y, c: z})) == smoment(robSol.subs({k: x, delta: y, c: z}), 4) + +def test_FiniteRV(): + F = FiniteRV('F', {1: S.Half, 2: Rational(1, 4), 3: Rational(1, 4)}, check=True) + p = Symbol("p", positive=True) + + assert dict(density(F).items()) == {S.One: S.Half, S(2): Rational(1, 4), S(3): Rational(1, 4)} + assert P(F >= 2) == S.Half + assert quantile(F)(p) == Piecewise((nan, p > S.One), (S.One, p <= S.Half),\ + (S(2), p <= Rational(3, 4)),(S(3), True)) + + assert pspace(F).domain.as_boolean() == Or( + *[Eq(F.symbol, i) for i in [1, 2, 3]]) + + assert F.pspace.domain.set == FiniteSet(1, 2, 3) + raises(ValueError, lambda: FiniteRV('F', {1: S.Half, 2: S.Half, 3: S.Half}, check=True)) + raises(ValueError, lambda: FiniteRV('F', {1: S.Half, 2: Rational(-1, 2), 3: S.One}, check=True)) + raises(ValueError, lambda: FiniteRV('F', {1: S.One, 2: Rational(3, 2), 3: S.Zero,\ + 4: Rational(-1, 2), 5: Rational(-3, 4), 6: Rational(-1, 4)}, check=True)) + + # purposeful invalid pmf but it should not raise since check=False + # see test_drv_types.test_ContinuousRV for explanation + X = FiniteRV('X', {1: 1, 2: 2}) + assert E(X) == 5 + assert P(X <= 2) + P(X > 2) != 1 + +def test_density_call(): + from sympy.abc import p + x = Bernoulli('x', p) + d = density(x) + assert d(0) == 1 - p + assert d(S.Zero) == 1 - p + assert d(5) == 0 + + assert 0 in d + assert 5 not in d + assert d(S.Zero) == d[S.Zero] + + +def test_DieDistribution(): + from sympy.abc import x + X = DieDistribution(6) + assert X.pmf(S.Half) is S.Zero + assert X.pmf(x).subs({x: 1}).doit() == Rational(1, 6) + assert X.pmf(x).subs({x: 7}).doit() == 0 + assert X.pmf(x).subs({x: -1}).doit() == 0 + assert X.pmf(x).subs({x: Rational(1, 3)}).doit() == 0 + raises(ValueError, lambda: X.pmf(Matrix([0, 0]))) + raises(ValueError, lambda: X.pmf(x**2 - 1)) + +def test_FinitePSpace(): + X = Die('X', 6) + space = pspace(X) + assert space.density == DieDistribution(6) + +def test_symbolic_conditions(): + B = Bernoulli('B', Rational(1, 4)) + D = Die('D', 4) + b, n = symbols('b, n') + Y = P(Eq(B, b)) + Z = E(D > n) + assert Y == \ + Piecewise((Rational(1, 4), Eq(b, 1)), (0, True)) + \ + Piecewise((Rational(3, 4), Eq(b, 0)), (0, True)) + assert Z == \ + Piecewise((Rational(1, 4), n < 1), (0, True)) + Piecewise((S.Half, n < 2), (0, True)) + \ + Piecewise((Rational(3, 4), n < 3), (0, True)) + Piecewise((S.One, n < 4), (0, True)) diff --git a/lib/python3.12/site-packages/sympy/stats/tests/test_joint_rv.py b/lib/python3.12/site-packages/sympy/stats/tests/test_joint_rv.py new file mode 100644 index 0000000000000000000000000000000000000000..057fc313dfbb31826b07fd1315205d22b86a7f96 --- /dev/null +++ b/lib/python3.12/site-packages/sympy/stats/tests/test_joint_rv.py @@ -0,0 +1,436 @@ +from sympy.concrete.products import Product +from sympy.concrete.summations import Sum +from sympy.core.numbers import (Rational, oo, pi) +from sympy.core.relational import Eq +from sympy.core.singleton import S +from sympy.core.symbol import symbols +from sympy.functions.combinatorial.factorials import (RisingFactorial, factorial) +from sympy.functions.elementary.complexes import polar_lift +from sympy.functions.elementary.exponential import exp +from sympy.functions.elementary.miscellaneous import sqrt +from sympy.functions.elementary.piecewise import Piecewise +from sympy.functions.special.bessel import besselk +from sympy.functions.special.gamma_functions import gamma +from sympy.matrices.dense import eye +from sympy.matrices.expressions.determinant import Determinant +from sympy.sets.fancysets import Range +from sympy.sets.sets import (Interval, ProductSet) +from sympy.simplify.simplify import simplify +from sympy.tensor.indexed import (Indexed, IndexedBase) +from sympy.core.numbers import comp +from sympy.integrals.integrals import integrate +from sympy.matrices import Matrix, MatrixSymbol +from sympy.matrices.expressions.matexpr import MatrixElement +from sympy.stats import density, median, marginal_distribution, Normal, Laplace, E, sample +from sympy.stats.joint_rv_types import (JointRV, MultivariateNormalDistribution, + JointDistributionHandmade, MultivariateT, NormalGamma, + GeneralizedMultivariateLogGammaOmega as GMVLGO, MultivariateBeta, + GeneralizedMultivariateLogGamma as GMVLG, MultivariateEwens, + Multinomial, NegativeMultinomial, MultivariateNormal, + MultivariateLaplace) +from sympy.testing.pytest import raises, XFAIL, skip, slow +from sympy.external import import_module + +from sympy.abc import x, y + + + +def test_Normal(): + m = Normal('A', [1, 2], [[1, 0], [0, 1]]) + A = MultivariateNormal('A', [1, 2], [[1, 0], [0, 1]]) + assert m == A + assert density(m)(1, 2) == 1/(2*pi) + assert m.pspace.distribution.set == ProductSet(S.Reals, S.Reals) + raises (ValueError, lambda:m[2]) + n = Normal('B', [1, 2, 3], [[1, 0, 0], [0, 1, 0], [0, 0, 1]]) + p = Normal('C', Matrix([1, 2]), Matrix([[1, 0], [0, 1]])) + assert density(m)(x, y) == density(p)(x, y) + assert marginal_distribution(n, 0, 1)(1, 2) == 1/(2*pi) + raises(ValueError, lambda: marginal_distribution(m)) + assert integrate(density(m)(x, y), (x, -oo, oo), (y, -oo, oo)).evalf() == 1.0 + N = Normal('N', [1, 2], [[x, 0], [0, y]]) + assert density(N)(0, 0) == exp(-((4*x + y)/(2*x*y)))/(2*pi*sqrt(x*y)) + + raises (ValueError, lambda: Normal('M', [1, 2], [[1, 1], [1, -1]])) + # symbolic + n = symbols('n', integer=True, positive=True) + mu = MatrixSymbol('mu', n, 1) + sigma = MatrixSymbol('sigma', n, n) + X = Normal('X', mu, sigma) + assert density(X) == MultivariateNormalDistribution(mu, sigma) + raises (NotImplementedError, lambda: median(m)) + # Below tests should work after issue #17267 is resolved + # assert E(X) == mu + # assert variance(X) == sigma + + # test symbolic multivariate normal densities + n = 3 + + Sg = MatrixSymbol('Sg', n, n) + mu = MatrixSymbol('mu', n, 1) + obs = MatrixSymbol('obs', n, 1) + + X = MultivariateNormal('X', mu, Sg) + density_X = density(X) + + eval_a = density_X(obs).subs({Sg: eye(3), + mu: Matrix([0, 0, 0]), obs: Matrix([0, 0, 0])}).doit() + eval_b = density_X(0, 0, 0).subs({Sg: eye(3), mu: Matrix([0, 0, 0])}).doit() + + assert eval_a == sqrt(2)/(4*pi**Rational(3/2)) + assert eval_b == sqrt(2)/(4*pi**Rational(3/2)) + + n = symbols('n', integer=True, positive=True) + + Sg = MatrixSymbol('Sg', n, n) + mu = MatrixSymbol('mu', n, 1) + obs = MatrixSymbol('obs', n, 1) + + X = MultivariateNormal('X', mu, Sg) + density_X_at_obs = density(X)(obs) + + expected_density = MatrixElement( + exp((S(1)/2) * (mu.T - obs.T) * Sg**(-1) * (-mu + obs)) / \ + sqrt((2*pi)**n * Determinant(Sg)), 0, 0) + + assert density_X_at_obs == expected_density + + +def test_MultivariateTDist(): + t1 = MultivariateT('T', [0, 0], [[1, 0], [0, 1]], 2) + assert(density(t1))(1, 1) == 1/(8*pi) + assert t1.pspace.distribution.set == ProductSet(S.Reals, S.Reals) + assert integrate(density(t1)(x, y), (x, -oo, oo), \ + (y, -oo, oo)).evalf() == 1.0 + raises(ValueError, lambda: MultivariateT('T', [1, 2], [[1, 1], [1, -1]], 1)) + t2 = MultivariateT('t2', [1, 2], [[x, 0], [0, y]], 1) + assert density(t2)(1, 2) == 1/(2*pi*sqrt(x*y)) + + +def test_multivariate_laplace(): + raises(ValueError, lambda: Laplace('T', [1, 2], [[1, 2], [2, 1]])) + L = Laplace('L', [1, 0], [[1, 0], [0, 1]]) + L2 = MultivariateLaplace('L2', [1, 0], [[1, 0], [0, 1]]) + assert density(L)(2, 3) == exp(2)*besselk(0, sqrt(39))/pi + L1 = Laplace('L1', [1, 2], [[x, 0], [0, y]]) + assert density(L1)(0, 1) == \ + exp(2/y)*besselk(0, sqrt((2 + 4/y + 1/x)/y))/(pi*sqrt(x*y)) + assert L.pspace.distribution.set == ProductSet(S.Reals, S.Reals) + assert L.pspace.distribution == L2.pspace.distribution + + +def test_NormalGamma(): + ng = NormalGamma('G', 1, 2, 3, 4) + assert density(ng)(1, 1) == 32*exp(-4)/sqrt(pi) + assert ng.pspace.distribution.set == ProductSet(S.Reals, Interval(0, oo)) + raises(ValueError, lambda:NormalGamma('G', 1, 2, 3, -1)) + assert marginal_distribution(ng, 0)(1) == \ + 3*sqrt(10)*gamma(Rational(7, 4))/(10*sqrt(pi)*gamma(Rational(5, 4))) + assert marginal_distribution(ng, y)(1) == exp(Rational(-1, 4))/128 + assert marginal_distribution(ng,[0,1])(x) == x**2*exp(-x/4)/128 + + +def test_GeneralizedMultivariateLogGammaDistribution(): + h = S.Half + omega = Matrix([[1, h, h, h], + [h, 1, h, h], + [h, h, 1, h], + [h, h, h, 1]]) + v, l, mu = (4, [1, 2, 3, 4], [1, 2, 3, 4]) + y_1, y_2, y_3, y_4 = symbols('y_1:5', real=True) + delta = symbols('d', positive=True) + G = GMVLGO('G', omega, v, l, mu) + Gd = GMVLG('Gd', delta, v, l, mu) + dend = ("d**4*Sum(4*24**(-n - 4)*(1 - d)**n*exp((n + 4)*(y_1 + 2*y_2 + 3*y_3 " + "+ 4*y_4) - exp(y_1) - exp(2*y_2)/2 - exp(3*y_3)/3 - exp(4*y_4)/4)/" + "(gamma(n + 1)*gamma(n + 4)**3), (n, 0, oo))") + assert str(density(Gd)(y_1, y_2, y_3, y_4)) == dend + den = ("5*2**(2/3)*5**(1/3)*Sum(4*24**(-n - 4)*(-2**(2/3)*5**(1/3)/4 + 1)**n*" + "exp((n + 4)*(y_1 + 2*y_2 + 3*y_3 + 4*y_4) - exp(y_1) - exp(2*y_2)/2 - " + "exp(3*y_3)/3 - exp(4*y_4)/4)/(gamma(n + 1)*gamma(n + 4)**3), (n, 0, oo))/64") + assert str(density(G)(y_1, y_2, y_3, y_4)) == den + marg = ("5*2**(2/3)*5**(1/3)*exp(4*y_1)*exp(-exp(y_1))*Integral(exp(-exp(4*G[3])" + "/4)*exp(16*G[3])*Integral(exp(-exp(3*G[2])/3)*exp(12*G[2])*Integral(exp(" + "-exp(2*G[1])/2)*exp(8*G[1])*Sum((-1/4)**n*(-4 + 2**(2/3)*5**(1/3" + "))**n*exp(n*y_1)*exp(2*n*G[1])*exp(3*n*G[2])*exp(4*n*G[3])/(24**n*gamma(n + 1)" + "*gamma(n + 4)**3), (n, 0, oo)), (G[1], -oo, oo)), (G[2], -oo, oo)), (G[3]" + ", -oo, oo))/5308416") + assert str(marginal_distribution(G, G[0])(y_1)) == marg + omega_f1 = Matrix([[1, h, h]]) + omega_f2 = Matrix([[1, h, h, h], + [h, 1, 2, h], + [h, h, 1, h], + [h, h, h, 1]]) + omega_f3 = Matrix([[6, h, h, h], + [h, 1, 2, h], + [h, h, 1, h], + [h, h, h, 1]]) + v_f = symbols("v_f", positive=False, real=True) + l_f = [1, 2, v_f, 4] + m_f = [v_f, 2, 3, 4] + omega_f4 = Matrix([[1, h, h, h, h], + [h, 1, h, h, h], + [h, h, 1, h, h], + [h, h, h, 1, h], + [h, h, h, h, 1]]) + l_f1 = [1, 2, 3, 4, 5] + omega_f5 = Matrix([[1]]) + mu_f5 = l_f5 = [1] + + raises(ValueError, lambda: GMVLGO('G', omega_f1, v, l, mu)) + raises(ValueError, lambda: GMVLGO('G', omega_f2, v, l, mu)) + raises(ValueError, lambda: GMVLGO('G', omega_f3, v, l, mu)) + raises(ValueError, lambda: GMVLGO('G', omega, v_f, l, mu)) + raises(ValueError, lambda: GMVLGO('G', omega, v, l_f, mu)) + raises(ValueError, lambda: GMVLGO('G', omega, v, l, m_f)) + raises(ValueError, lambda: GMVLGO('G', omega_f4, v, l, mu)) + raises(ValueError, lambda: GMVLGO('G', omega, v, l_f1, mu)) + raises(ValueError, lambda: GMVLGO('G', omega_f5, v, l_f5, mu_f5)) + raises(ValueError, lambda: GMVLG('G', Rational(3, 2), v, l, mu)) + + +def test_MultivariateBeta(): + a1, a2 = symbols('a1, a2', positive=True) + a1_f, a2_f = symbols('a1, a2', positive=False, real=True) + mb = MultivariateBeta('B', [a1, a2]) + mb_c = MultivariateBeta('C', a1, a2) + assert density(mb)(1, 2) == S(2)**(a2 - 1)*gamma(a1 + a2)/\ + (gamma(a1)*gamma(a2)) + assert marginal_distribution(mb_c, 0)(3) == S(3)**(a1 - 1)*gamma(a1 + a2)/\ + (a2*gamma(a1)*gamma(a2)) + raises(ValueError, lambda: MultivariateBeta('b1', [a1_f, a2])) + raises(ValueError, lambda: MultivariateBeta('b2', [a1, a2_f])) + raises(ValueError, lambda: MultivariateBeta('b3', [0, 0])) + raises(ValueError, lambda: MultivariateBeta('b4', [a1_f, a2_f])) + assert mb.pspace.distribution.set == ProductSet(Interval(0, 1), Interval(0, 1)) + + +def test_MultivariateEwens(): + n, theta, i = symbols('n theta i', positive=True) + + # tests for integer dimensions + theta_f = symbols('t_f', negative=True) + a = symbols('a_1:4', positive = True, integer = True) + ed = MultivariateEwens('E', 3, theta) + assert density(ed)(a[0], a[1], a[2]) == Piecewise((6*2**(-a[1])*3**(-a[2])* + theta**a[0]*theta**a[1]*theta**a[2]/ + (theta*(theta + 1)*(theta + 2)* + factorial(a[0])*factorial(a[1])* + factorial(a[2])), Eq(a[0] + 2*a[1] + + 3*a[2], 3)), (0, True)) + assert marginal_distribution(ed, ed[1])(a[1]) == Piecewise((6*2**(-a[1])* + theta**a[1]/((theta + 1)* + (theta + 2)*factorial(a[1])), + Eq(2*a[1] + 1, 3)), (0, True)) + raises(ValueError, lambda: MultivariateEwens('e1', 5, theta_f)) + assert ed.pspace.distribution.set == ProductSet(Range(0, 4, 1), + Range(0, 2, 1), Range(0, 2, 1)) + + # tests for symbolic dimensions + eds = MultivariateEwens('E', n, theta) + a = IndexedBase('a') + j, k = symbols('j, k') + den = Piecewise((factorial(n)*Product(theta**a[j]*(j + 1)**(-a[j])/ + factorial(a[j]), (j, 0, n - 1))/RisingFactorial(theta, n), + Eq(n, Sum((k + 1)*a[k], (k, 0, n - 1)))), (0, True)) + assert density(eds)(a).dummy_eq(den) + + +def test_Multinomial(): + n, x1, x2, x3, x4 = symbols('n, x1, x2, x3, x4', nonnegative=True, integer=True) + p1, p2, p3, p4 = symbols('p1, p2, p3, p4', positive=True) + p1_f, n_f = symbols('p1_f, n_f', negative=True) + M = Multinomial('M', n, [p1, p2, p3, p4]) + C = Multinomial('C', 3, p1, p2, p3) + f = factorial + assert density(M)(x1, x2, x3, x4) == Piecewise((p1**x1*p2**x2*p3**x3*p4**x4* + f(n)/(f(x1)*f(x2)*f(x3)*f(x4)), + Eq(n, x1 + x2 + x3 + x4)), (0, True)) + assert marginal_distribution(C, C[0])(x1).subs(x1, 1) ==\ + 3*p1*p2**2 +\ + 6*p1*p2*p3 +\ + 3*p1*p3**2 + raises(ValueError, lambda: Multinomial('b1', 5, [p1, p2, p3, p1_f])) + raises(ValueError, lambda: Multinomial('b2', n_f, [p1, p2, p3, p4])) + raises(ValueError, lambda: Multinomial('b3', n, 0.5, 0.4, 0.3, 0.1)) + + +def test_NegativeMultinomial(): + k0, x1, x2, x3, x4 = symbols('k0, x1, x2, x3, x4', nonnegative=True, integer=True) + p1, p2, p3, p4 = symbols('p1, p2, p3, p4', positive=True) + p1_f = symbols('p1_f', negative=True) + N = NegativeMultinomial('N', 4, [p1, p2, p3, p4]) + C = NegativeMultinomial('C', 4, 0.1, 0.2, 0.3) + g = gamma + f = factorial + assert simplify(density(N)(x1, x2, x3, x4) - + p1**x1*p2**x2*p3**x3*p4**x4*(-p1 - p2 - p3 - p4 + 1)**4*g(x1 + x2 + + x3 + x4 + 4)/(6*f(x1)*f(x2)*f(x3)*f(x4))) is S.Zero + assert comp(marginal_distribution(C, C[0])(1).evalf(), 0.33, .01) + raises(ValueError, lambda: NegativeMultinomial('b1', 5, [p1, p2, p3, p1_f])) + raises(ValueError, lambda: NegativeMultinomial('b2', k0, 0.5, 0.4, 0.3, 0.4)) + assert N.pspace.distribution.set == ProductSet(Range(0, oo, 1), + Range(0, oo, 1), Range(0, oo, 1), Range(0, oo, 1)) + + +@slow +def test_JointPSpace_marginal_distribution(): + T = MultivariateT('T', [0, 0], [[1, 0], [0, 1]], 2) + got = marginal_distribution(T, T[1])(x) + ans = sqrt(2)*(x**2/2 + 1)/(4*polar_lift(x**2/2 + 1)**(S(5)/2)) + assert got == ans, got + assert integrate(marginal_distribution(T, 1)(x), (x, -oo, oo)) == 1 + + t = MultivariateT('T', [0, 0, 0], [[1, 0, 0], [0, 1, 0], [0, 0, 1]], 3) + assert comp(marginal_distribution(t, 0)(1).evalf(), 0.2, .01) + + +def test_JointRV(): + x1, x2 = (Indexed('x', i) for i in (1, 2)) + pdf = exp(-x1**2/2 + x1 - x2**2/2 - S.Half)/(2*pi) + X = JointRV('x', pdf) + assert density(X)(1, 2) == exp(-2)/(2*pi) + assert isinstance(X.pspace.distribution, JointDistributionHandmade) + assert marginal_distribution(X, 0)(2) == sqrt(2)*exp(Rational(-1, 2))/(2*sqrt(pi)) + + +def test_expectation(): + m = Normal('A', [x, y], [[1, 0], [0, 1]]) + assert simplify(E(m[1])) == y + + +@XFAIL +def test_joint_vector_expectation(): + m = Normal('A', [x, y], [[1, 0], [0, 1]]) + assert E(m) == (x, y) + + +def test_sample_numpy(): + distribs_numpy = [ + MultivariateNormal("M", [3, 4], [[2, 1], [1, 2]]), + MultivariateBeta("B", [0.4, 5, 15, 50, 203]), + Multinomial("N", 50, [0.3, 0.2, 0.1, 0.25, 0.15]) + ] + size = 3 + numpy = import_module('numpy') + if not numpy: + skip('Numpy is not installed. Abort tests for _sample_numpy.') + else: + for X in distribs_numpy: + samps = sample(X, size=size, library='numpy') + for sam in samps: + assert tuple(sam) in X.pspace.distribution.set + N_c = NegativeMultinomial('N', 3, 0.1, 0.1, 0.1) + raises(NotImplementedError, lambda: sample(N_c, library='numpy')) + + +def test_sample_scipy(): + distribs_scipy = [ + MultivariateNormal("M", [0, 0], [[0.1, 0.025], [0.025, 0.1]]), + MultivariateBeta("B", [0.4, 5, 15]), + Multinomial("N", 8, [0.3, 0.2, 0.1, 0.4]) + ] + + size = 3 + scipy = import_module('scipy') + if not scipy: + skip('Scipy not installed. Abort tests for _sample_scipy.') + else: + for X in distribs_scipy: + samps = sample(X, size=size) + samps2 = sample(X, size=(2, 2)) + for sam in samps: + assert tuple(sam) in X.pspace.distribution.set + for i in range(2): + for j in range(2): + assert tuple(samps2[i][j]) in X.pspace.distribution.set + N_c = NegativeMultinomial('N', 3, 0.1, 0.1, 0.1) + raises(NotImplementedError, lambda: sample(N_c)) + + +def test_sample_pymc(): + distribs_pymc = [ + MultivariateNormal("M", [5, 2], [[1, 0], [0, 1]]), + MultivariateBeta("B", [0.4, 5, 15]), + Multinomial("N", 4, [0.3, 0.2, 0.1, 0.4]) + ] + size = 3 + pymc = import_module('pymc') + if not pymc: + skip('PyMC is not installed. Abort tests for _sample_pymc.') + else: + for X in distribs_pymc: + samps = sample(X, size=size, library='pymc') + for sam in samps: + assert tuple(sam.flatten()) in X.pspace.distribution.set + N_c = NegativeMultinomial('N', 3, 0.1, 0.1, 0.1) + raises(NotImplementedError, lambda: sample(N_c, library='pymc')) + + +def test_sample_seed(): + x1, x2 = (Indexed('x', i) for i in (1, 2)) + pdf = exp(-x1**2/2 + x1 - x2**2/2 - S.Half)/(2*pi) + X = JointRV('x', pdf) + + libraries = ['scipy', 'numpy', 'pymc'] + for lib in libraries: + try: + imported_lib = import_module(lib) + if imported_lib: + s0, s1, s2 = [], [], [] + s0 = sample(X, size=10, library=lib, seed=0) + s1 = sample(X, size=10, library=lib, seed=0) + s2 = sample(X, size=10, library=lib, seed=1) + assert all(s0 == s1) + assert all(s1 != s2) + except NotImplementedError: + continue + +# +# XXX: This fails for pymc. Previously the test appeared to pass but that is +# just because the library argument was not passed so the test always used +# scipy. +# +def test_issue_21057(): + m = Normal("x", [0, 0], [[0, 0], [0, 0]]) + n = MultivariateNormal("x", [0, 0], [[0, 0], [0, 0]]) + p = Normal("x", [0, 0], [[0, 0], [0, 1]]) + assert m == n + libraries = ('scipy', 'numpy') # , 'pymc') # <-- pymc fails + for library in libraries: + try: + imported_lib = import_module(library) + if imported_lib: + s1 = sample(m, size=8, library=library) + s2 = sample(n, size=8, library=library) + s3 = sample(p, size=8, library=library) + assert tuple(s1.flatten()) == tuple(s2.flatten()) + for s in s3: + assert tuple(s.flatten()) in p.pspace.distribution.set + except NotImplementedError: + continue + + +# +# When this passes the pymc part can be uncommented in test_issue_21057 above +# and this can be deleted. +# +@XFAIL +def test_issue_21057_pymc(): + m = Normal("x", [0, 0], [[0, 0], [0, 0]]) + n = MultivariateNormal("x", [0, 0], [[0, 0], [0, 0]]) + p = Normal("x", [0, 0], [[0, 0], [0, 1]]) + assert m == n + libraries = ('pymc',) + for library in libraries: + try: + imported_lib = import_module(library) + if imported_lib: + s1 = sample(m, size=8, library=library) + s2 = sample(n, size=8, library=library) + s3 = sample(p, size=8, library=library) + assert tuple(s1.flatten()) == tuple(s2.flatten()) + for s in s3: + assert tuple(s.flatten()) in p.pspace.distribution.set + except NotImplementedError: + continue diff --git a/lib/python3.12/site-packages/sympy/stats/tests/test_matrix_distributions.py b/lib/python3.12/site-packages/sympy/stats/tests/test_matrix_distributions.py new file mode 100644 index 0000000000000000000000000000000000000000..a2a2dcdd853793d9f77e1a88adf63158ed68e3ba --- /dev/null +++ b/lib/python3.12/site-packages/sympy/stats/tests/test_matrix_distributions.py @@ -0,0 +1,186 @@ +from sympy.concrete.products import Product +from sympy.core.numbers import pi +from sympy.core.singleton import S +from sympy.core.symbol import (Dummy, symbols) +from sympy.functions.elementary.exponential import exp +from sympy.functions.elementary.miscellaneous import sqrt +from sympy.functions.special.gamma_functions import gamma +from sympy.matrices import Determinant, Matrix, Trace, MatrixSymbol, MatrixSet +from sympy.stats import density, sample +from sympy.stats.matrix_distributions import (MatrixGammaDistribution, + MatrixGamma, MatrixPSpace, Wishart, MatrixNormal, MatrixStudentT) +from sympy.testing.pytest import raises, skip +from sympy.external import import_module + + +def test_MatrixPSpace(): + M = MatrixGammaDistribution(1, 2, [[2, 1], [1, 2]]) + MP = MatrixPSpace('M', M, 2, 2) + assert MP.distribution == M + raises(ValueError, lambda: MatrixPSpace('M', M, 1.2, 2)) + +def test_MatrixGamma(): + M = MatrixGamma('M', 1, 2, [[1, 0], [0, 1]]) + assert M.pspace.distribution.set == MatrixSet(2, 2, S.Reals) + assert isinstance(density(M), MatrixGammaDistribution) + X = MatrixSymbol('X', 2, 2) + num = exp(Trace(Matrix([[-S(1)/2, 0], [0, -S(1)/2]])*X)) + assert density(M)(X).doit() == num/(4*pi*sqrt(Determinant(X))) + assert density(M)([[2, 1], [1, 2]]).doit() == sqrt(3)*exp(-2)/(12*pi) + X = MatrixSymbol('X', 1, 2) + Y = MatrixSymbol('Y', 1, 2) + assert density(M)([X, Y]).doit() == exp(-X[0, 0]/2 - Y[0, 1]/2)/(4*pi*sqrt( + X[0, 0]*Y[0, 1] - X[0, 1]*Y[0, 0])) + # symbolic + a, b = symbols('a b', positive=True) + d = symbols('d', positive=True, integer=True) + Y = MatrixSymbol('Y', d, d) + Z = MatrixSymbol('Z', 2, 2) + SM = MatrixSymbol('SM', d, d) + M2 = MatrixGamma('M2', a, b, SM) + M3 = MatrixGamma('M3', 2, 3, [[2, 1], [1, 2]]) + k = Dummy('k') + exprd = pi**(-d*(d - 1)/4)*b**(-a*d)*exp(Trace((-1/b)*SM**(-1)*Y) + )*Determinant(SM)**(-a)*Determinant(Y)**(a - d/2 - S(1)/2)/Product( + gamma(-k/2 + a + S(1)/2), (k, 1, d)) + assert density(M2)(Y).dummy_eq(exprd) + raises(NotImplementedError, lambda: density(M3 + M)(Z)) + raises(ValueError, lambda: density(M)(1)) + raises(ValueError, lambda: MatrixGamma('M', -1, 2, [[1, 0], [0, 1]])) + raises(ValueError, lambda: MatrixGamma('M', -1, -2, [[1, 0], [0, 1]])) + raises(ValueError, lambda: MatrixGamma('M', -1, 2, [[1, 0], [2, 1]])) + raises(ValueError, lambda: MatrixGamma('M', -1, 2, [[1, 0], [0]])) + +def test_Wishart(): + W = Wishart('W', 5, [[1, 0], [0, 1]]) + assert W.pspace.distribution.set == MatrixSet(2, 2, S.Reals) + X = MatrixSymbol('X', 2, 2) + term1 = exp(Trace(Matrix([[-S(1)/2, 0], [0, -S(1)/2]])*X)) + assert density(W)(X).doit() == term1 * Determinant(X)/(24*pi) + assert density(W)([[2, 1], [1, 2]]).doit() == exp(-2)/(8*pi) + n = symbols('n', positive=True) + d = symbols('d', positive=True, integer=True) + Y = MatrixSymbol('Y', d, d) + SM = MatrixSymbol('SM', d, d) + W = Wishart('W', n, SM) + k = Dummy('k') + exprd = 2**(-d*n/2)*pi**(-d*(d - 1)/4)*exp(Trace(-(S(1)/2)*SM**(-1)*Y) + )*Determinant(SM)**(-n/2)*Determinant(Y)**( + -d/2 + n/2 - S(1)/2)/Product(gamma(-k/2 + n/2 + S(1)/2), (k, 1, d)) + assert density(W)(Y).dummy_eq(exprd) + raises(ValueError, lambda: density(W)(1)) + raises(ValueError, lambda: Wishart('W', -1, [[1, 0], [0, 1]])) + raises(ValueError, lambda: Wishart('W', -1, [[1, 0], [2, 1]])) + raises(ValueError, lambda: Wishart('W', 2, [[1, 0], [0]])) + +def test_MatrixNormal(): + M = MatrixNormal('M', [[5, 6]], [4], [[2, 1], [1, 2]]) + assert M.pspace.distribution.set == MatrixSet(1, 2, S.Reals) + X = MatrixSymbol('X', 1, 2) + term1 = exp(-Trace(Matrix([[ S(2)/3, -S(1)/3], [-S(1)/3, S(2)/3]])*( + Matrix([[-5], [-6]]) + X.T)*Matrix([[S(1)/4]])*(Matrix([[-5, -6]]) + X))/2) + assert density(M)(X).doit() == (sqrt(3)) * term1/(24*pi) + assert density(M)([[7, 8]]).doit() == sqrt(3)*exp(-S(1)/3)/(24*pi) + d, n = symbols('d n', positive=True, integer=True) + SM2 = MatrixSymbol('SM2', d, d) + SM1 = MatrixSymbol('SM1', n, n) + LM = MatrixSymbol('LM', n, d) + Y = MatrixSymbol('Y', n, d) + M = MatrixNormal('M', LM, SM1, SM2) + exprd = (2*pi)**(-d*n/2)*exp(-Trace(SM2**(-1)*(-LM.T + Y.T)*SM1**(-1)*(-LM + Y) + )/2)*Determinant(SM1)**(-d/2)*Determinant(SM2)**(-n/2) + assert density(M)(Y).doit() == exprd + raises(ValueError, lambda: density(M)(1)) + raises(ValueError, lambda: MatrixNormal('M', [1, 2], [[1, 0], [0, 1]], [[1, 0], [2, 1]])) + raises(ValueError, lambda: MatrixNormal('M', [1, 2], [[1, 0], [2, 1]], [[1, 0], [0, 1]])) + raises(ValueError, lambda: MatrixNormal('M', [1, 2], [[1, 0], [0, 1]], [[1, 0], [0, 1]])) + raises(ValueError, lambda: MatrixNormal('M', [1, 2], [[1, 0], [2]], [[1, 0], [0, 1]])) + raises(ValueError, lambda: MatrixNormal('M', [1, 2], [[1, 0], [2, 1]], [[1, 0], [0]])) + raises(ValueError, lambda: MatrixNormal('M', [[1, 2]], [[1, 0], [0, 1]], [[1, 0]])) + raises(ValueError, lambda: MatrixNormal('M', [[1, 2]], [1], [[1, 0]])) + +def test_MatrixStudentT(): + M = MatrixStudentT('M', 2, [[5, 6]], [[2, 1], [1, 2]], [4]) + assert M.pspace.distribution.set == MatrixSet(1, 2, S.Reals) + X = MatrixSymbol('X', 1, 2) + D = pi ** (-1.0) * Determinant(Matrix([[4]])) ** (-1.0) * Determinant(Matrix([[2, 1], [1, 2]])) \ + ** (-0.5) / Determinant(Matrix([[S(1) / 4]]) * (Matrix([[-5, -6]]) + X) + * Matrix([[S(2) / 3, -S(1) / 3], [-S(1) / 3, S(2) / 3]]) * ( + Matrix([[-5], [-6]]) + X.T) + Matrix([[1]])) ** 2 + assert density(M)(X) == D + + v = symbols('v', positive=True) + n, p = 1, 2 + Omega = MatrixSymbol('Omega', p, p) + Sigma = MatrixSymbol('Sigma', n, n) + Location = MatrixSymbol('Location', n, p) + Y = MatrixSymbol('Y', n, p) + M = MatrixStudentT('M', v, Location, Omega, Sigma) + + exprd = gamma(v/2 + 1)*Determinant(Matrix([[1]]) + Sigma**(-1)*(-Location + Y)*Omega**(-1)*(-Location.T + Y.T))**(-v/2 - 1) / \ + (pi*gamma(v/2)*sqrt(Determinant(Omega))*Determinant(Sigma)) + + assert density(M)(Y) == exprd + raises(ValueError, lambda: density(M)(1)) + raises(ValueError, lambda: MatrixStudentT('M', 1, [1, 2], [[1, 0], [0, 1]], [[1, 0], [2, 1]])) + raises(ValueError, lambda: MatrixStudentT('M', 1, [1, 2], [[1, 0], [2, 1]], [[1, 0], [0, 1]])) + raises(ValueError, lambda: MatrixStudentT('M', 1, [1, 2], [[1, 0], [0, 1]], [[1, 0], [0, 1]])) + raises(ValueError, lambda: MatrixStudentT('M', 1, [1, 2], [[1, 0], [2]], [[1, 0], [0, 1]])) + raises(ValueError, lambda: MatrixStudentT('M', 1, [1, 2], [[1, 0], [2, 1]], [[1], [2]])) + raises(ValueError, lambda: MatrixStudentT('M', 1, [[1, 2]], [[1, 0], [0, 1]], [[1, 0]])) + raises(ValueError, lambda: MatrixStudentT('M', 1, [[1, 2]], [1], [[1, 0]])) + raises(ValueError, lambda: MatrixStudentT('M', -1, [1, 2], [[1, 0], [0, 1]], [4])) + +def test_sample_scipy(): + distribs_scipy = [ + MatrixNormal('M', [[5, 6]], [4], [[2, 1], [1, 2]]), + Wishart('W', 5, [[1, 0], [0, 1]]) + ] + + size = 5 + scipy = import_module('scipy') + if not scipy: + skip('Scipy not installed. Abort tests for _sample_scipy.') + else: + for X in distribs_scipy: + samps = sample(X, size=size) + for sam in samps: + assert Matrix(sam) in X.pspace.distribution.set + M = MatrixGamma('M', 1, 2, [[1, 0], [0, 1]]) + raises(NotImplementedError, lambda: sample(M, size=3)) + +def test_sample_pymc(): + distribs_pymc = [ + MatrixNormal('M', [[5, 6], [3, 4]], [[1, 0], [0, 1]], [[2, 1], [1, 2]]), + Wishart('W', 7, [[2, 1], [1, 2]]) + ] + size = 3 + pymc = import_module('pymc') + if not pymc: + skip('PyMC is not installed. Abort tests for _sample_pymc.') + else: + for X in distribs_pymc: + samps = sample(X, size=size, library='pymc') + for sam in samps: + assert Matrix(sam) in X.pspace.distribution.set + M = MatrixGamma('M', 1, 2, [[1, 0], [0, 1]]) + raises(NotImplementedError, lambda: sample(M, size=3)) + +def test_sample_seed(): + X = MatrixNormal('M', [[5, 6], [3, 4]], [[1, 0], [0, 1]], [[2, 1], [1, 2]]) + + libraries = ['scipy', 'numpy', 'pymc'] + for lib in libraries: + try: + imported_lib = import_module(lib) + if imported_lib: + s0, s1, s2 = [], [], [] + s0 = sample(X, size=10, library=lib, seed=0) + s1 = sample(X, size=10, library=lib, seed=0) + s2 = sample(X, size=10, library=lib, seed=1) + for i in range(10): + assert (s0[i] == s1[i]).all() + assert (s1[i] != s2[i]).all() + + except NotImplementedError: + continue diff --git a/lib/python3.12/site-packages/sympy/stats/tests/test_mix.py b/lib/python3.12/site-packages/sympy/stats/tests/test_mix.py new file mode 100644 index 0000000000000000000000000000000000000000..4334d9b144a5ddaad938f195f0276e0e8993aa35 --- /dev/null +++ b/lib/python3.12/site-packages/sympy/stats/tests/test_mix.py @@ -0,0 +1,82 @@ +from sympy.concrete.summations import Sum +from sympy.core.add import Add +from sympy.core.mul import Mul +from sympy.core.numbers import (Integer, oo, pi) +from sympy.core.power import Pow +from sympy.core.relational import (Eq, Ne) +from sympy.core.symbol import (Dummy, Symbol, symbols) +from sympy.functions.combinatorial.factorials import factorial +from sympy.functions.elementary.exponential import exp +from sympy.functions.elementary.miscellaneous import sqrt +from sympy.functions.elementary.piecewise import Piecewise +from sympy.functions.special.delta_functions import DiracDelta +from sympy.functions.special.gamma_functions import gamma +from sympy.integrals.integrals import Integral +from sympy.simplify.simplify import simplify +from sympy.tensor.indexed import (Indexed, IndexedBase) +from sympy.functions.elementary.piecewise import ExprCondPair +from sympy.stats import (Poisson, Beta, Exponential, P, + Multinomial, MultivariateBeta) +from sympy.stats.crv_types import Normal +from sympy.stats.drv_types import PoissonDistribution +from sympy.stats.compound_rv import CompoundPSpace, CompoundDistribution +from sympy.stats.joint_rv import MarginalDistribution +from sympy.stats.rv import pspace, density +from sympy.testing.pytest import ignore_warnings + +def test_density(): + x = Symbol('x') + l = Symbol('l', positive=True) + rate = Beta(l, 2, 3) + X = Poisson(x, rate) + assert isinstance(pspace(X), CompoundPSpace) + assert density(X, Eq(rate, rate.symbol)) == PoissonDistribution(l) + N1 = Normal('N1', 0, 1) + N2 = Normal('N2', N1, 2) + assert density(N2)(0).doit() == sqrt(10)/(10*sqrt(pi)) + assert simplify(density(N2, Eq(N1, 1))(x)) == \ + sqrt(2)*exp(-(x - 1)**2/8)/(4*sqrt(pi)) + assert simplify(density(N2)(x)) == sqrt(10)*exp(-x**2/10)/(10*sqrt(pi)) + +def test_MarginalDistribution(): + a1, p1, p2 = symbols('a1 p1 p2', positive=True) + C = Multinomial('C', 2, p1, p2) + B = MultivariateBeta('B', a1, C[0]) + MGR = MarginalDistribution(B, (C[0],)) + mgrc = Mul(Symbol('B'), Piecewise(ExprCondPair(Mul(Integer(2), + Pow(Symbol('p1', positive=True), Indexed(IndexedBase(Symbol('C')), + Integer(0))), Pow(Symbol('p2', positive=True), + Indexed(IndexedBase(Symbol('C')), Integer(1))), + Pow(factorial(Indexed(IndexedBase(Symbol('C')), Integer(0))), Integer(-1)), + Pow(factorial(Indexed(IndexedBase(Symbol('C')), Integer(1))), Integer(-1))), + Eq(Add(Indexed(IndexedBase(Symbol('C')), Integer(0)), + Indexed(IndexedBase(Symbol('C')), Integer(1))), Integer(2))), + ExprCondPair(Integer(0), True)), Pow(gamma(Symbol('a1', positive=True)), + Integer(-1)), gamma(Add(Symbol('a1', positive=True), + Indexed(IndexedBase(Symbol('C')), Integer(0)))), + Pow(gamma(Indexed(IndexedBase(Symbol('C')), Integer(0))), Integer(-1)), + Pow(Indexed(IndexedBase(Symbol('B')), Integer(0)), + Add(Symbol('a1', positive=True), Integer(-1))), + Pow(Indexed(IndexedBase(Symbol('B')), Integer(1)), + Add(Indexed(IndexedBase(Symbol('C')), Integer(0)), Integer(-1)))) + assert MGR(C) == mgrc + +def test_compound_distribution(): + Y = Poisson('Y', 1) + Z = Poisson('Z', Y) + assert isinstance(pspace(Z), CompoundPSpace) + assert isinstance(pspace(Z).distribution, CompoundDistribution) + assert Z.pspace.distribution.pdf(1).doit() == exp(-2)*exp(exp(-1)) + +def test_mix_expression(): + Y, E = Poisson('Y', 1), Exponential('E', 1) + k = Dummy('k') + expr1 = Integral(Sum(exp(-1)*Integral(exp(-k)*DiracDelta(k - 2), (k, 0, oo) + )/factorial(k), (k, 0, oo)), (k, -oo, 0)) + expr2 = Integral(Sum(exp(-1)*Integral(exp(-k)*DiracDelta(k - 2), (k, 0, oo) + )/factorial(k), (k, 0, oo)), (k, 0, oo)) + assert P(Eq(Y + E, 1)) == 0 + assert P(Ne(Y + E, 2)) == 1 + with ignore_warnings(UserWarning): ### TODO: Restore tests once warnings are removed + assert P(E + Y < 2, evaluate=False).rewrite(Integral).dummy_eq(expr1) + assert P(E + Y > 2, evaluate=False).rewrite(Integral).dummy_eq(expr2) diff --git a/lib/python3.12/site-packages/sympy/stats/tests/test_random_matrix.py b/lib/python3.12/site-packages/sympy/stats/tests/test_random_matrix.py new file mode 100644 index 0000000000000000000000000000000000000000..ba570a16bc42620d53bce19be71e7d125965ede1 --- /dev/null +++ b/lib/python3.12/site-packages/sympy/stats/tests/test_random_matrix.py @@ -0,0 +1,135 @@ +from sympy.concrete.products import Product +from sympy.core.function import Lambda +from sympy.core.numbers import (I, Rational, pi) +from sympy.core.singleton import S +from sympy.core.symbol import Dummy +from sympy.functions.elementary.complexes import Abs +from sympy.functions.elementary.exponential import exp +from sympy.functions.elementary.miscellaneous import sqrt +from sympy.integrals.integrals import Integral +from sympy.matrices.dense import Matrix +from sympy.matrices.expressions.matexpr import MatrixSymbol +from sympy.matrices.expressions.trace import Trace +from sympy.tensor.indexed import IndexedBase +from sympy.stats import (GaussianUnitaryEnsemble as GUE, density, + GaussianOrthogonalEnsemble as GOE, + GaussianSymplecticEnsemble as GSE, + joint_eigen_distribution, + CircularUnitaryEnsemble as CUE, + CircularOrthogonalEnsemble as COE, + CircularSymplecticEnsemble as CSE, + JointEigenDistribution, + level_spacing_distribution, + Normal, Beta) +from sympy.stats.joint_rv_types import JointDistributionHandmade +from sympy.stats.rv import RandomMatrixSymbol +from sympy.stats.random_matrix_models import GaussianEnsemble, RandomMatrixPSpace +from sympy.testing.pytest import raises + +def test_GaussianEnsemble(): + G = GaussianEnsemble('G', 3) + assert density(G) == G.pspace.model + raises(ValueError, lambda: GaussianEnsemble('G', 3.5)) + +def test_GaussianUnitaryEnsemble(): + H = RandomMatrixSymbol('H', 3, 3) + G = GUE('U', 3) + assert density(G)(H) == sqrt(2)*exp(-3*Trace(H**2)/2)/(4*pi**Rational(9, 2)) + i, j = (Dummy('i', integer=True, positive=True), + Dummy('j', integer=True, positive=True)) + l = IndexedBase('l') + assert joint_eigen_distribution(G).dummy_eq( + Lambda((l[1], l[2], l[3]), + 27*sqrt(6)*exp(-3*(l[1]**2)/2 - 3*(l[2]**2)/2 - 3*(l[3]**2)/2)* + Product(Abs(l[i] - l[j])**2, (j, i + 1, 3), (i, 1, 2))/(16*pi**Rational(3, 2)))) + s = Dummy('s') + assert level_spacing_distribution(G).dummy_eq(Lambda(s, 32*s**2*exp(-4*s**2/pi)/pi**2)) + + +def test_GaussianOrthogonalEnsemble(): + H = RandomMatrixSymbol('H', 3, 3) + _H = MatrixSymbol('_H', 3, 3) + G = GOE('O', 3) + assert density(G)(H) == exp(-3*Trace(H**2)/4)/Integral(exp(-3*Trace(_H**2)/4), _H) + i, j = (Dummy('i', integer=True, positive=True), + Dummy('j', integer=True, positive=True)) + l = IndexedBase('l') + assert joint_eigen_distribution(G).dummy_eq( + Lambda((l[1], l[2], l[3]), + 9*sqrt(2)*exp(-3*l[1]**2/2 - 3*l[2]**2/2 - 3*l[3]**2/2)* + Product(Abs(l[i] - l[j]), (j, i + 1, 3), (i, 1, 2))/(32*pi))) + s = Dummy('s') + assert level_spacing_distribution(G).dummy_eq(Lambda(s, s*pi*exp(-s**2*pi/4)/2)) + +def test_GaussianSymplecticEnsemble(): + H = RandomMatrixSymbol('H', 3, 3) + _H = MatrixSymbol('_H', 3, 3) + G = GSE('O', 3) + assert density(G)(H) == exp(-3*Trace(H**2))/Integral(exp(-3*Trace(_H**2)), _H) + i, j = (Dummy('i', integer=True, positive=True), + Dummy('j', integer=True, positive=True)) + l = IndexedBase('l') + assert joint_eigen_distribution(G).dummy_eq( + Lambda((l[1], l[2], l[3]), + 162*sqrt(3)*exp(-3*l[1]**2/2 - 3*l[2]**2/2 - 3*l[3]**2/2)* + Product(Abs(l[i] - l[j])**4, (j, i + 1, 3), (i, 1, 2))/(5*pi**Rational(3, 2)))) + s = Dummy('s') + assert level_spacing_distribution(G).dummy_eq(Lambda(s, S(262144)*s**4*exp(-64*s**2/(9*pi))/(729*pi**3))) + +def test_CircularUnitaryEnsemble(): + CU = CUE('U', 3) + j, k = (Dummy('j', integer=True, positive=True), + Dummy('k', integer=True, positive=True)) + t = IndexedBase('t') + assert joint_eigen_distribution(CU).dummy_eq( + Lambda((t[1], t[2], t[3]), + Product(Abs(exp(I*t[j]) - exp(I*t[k]))**2, + (j, k + 1, 3), (k, 1, 2))/(48*pi**3)) + ) + +def test_CircularOrthogonalEnsemble(): + CO = COE('U', 3) + j, k = (Dummy('j', integer=True, positive=True), + Dummy('k', integer=True, positive=True)) + t = IndexedBase('t') + assert joint_eigen_distribution(CO).dummy_eq( + Lambda((t[1], t[2], t[3]), + Product(Abs(exp(I*t[j]) - exp(I*t[k])), + (j, k + 1, 3), (k, 1, 2))/(48*pi**2)) + ) + +def test_CircularSymplecticEnsemble(): + CS = CSE('U', 3) + j, k = (Dummy('j', integer=True, positive=True), + Dummy('k', integer=True, positive=True)) + t = IndexedBase('t') + assert joint_eigen_distribution(CS).dummy_eq( + Lambda((t[1], t[2], t[3]), + Product(Abs(exp(I*t[j]) - exp(I*t[k]))**4, + (j, k + 1, 3), (k, 1, 2))/(720*pi**3)) + ) + +def test_JointEigenDistribution(): + A = Matrix([[Normal('A00', 0, 1), Normal('A01', 1, 1)], + [Beta('A10', 1, 1), Beta('A11', 1, 1)]]) + assert JointEigenDistribution(A) == \ + JointDistributionHandmade(-sqrt(A[0, 0]**2 - 2*A[0, 0]*A[1, 1] + 4*A[0, 1]*A[1, 0] + A[1, 1]**2)/2 + + A[0, 0]/2 + A[1, 1]/2, sqrt(A[0, 0]**2 - 2*A[0, 0]*A[1, 1] + 4*A[0, 1]*A[1, 0] + A[1, 1]**2)/2 + A[0, 0]/2 + A[1, 1]/2) + raises(ValueError, lambda: JointEigenDistribution(Matrix([[1, 0], [2, 1]]))) + +def test_issue_19841(): + G1 = GUE('U', 2) + G2 = G1.xreplace({2: 2}) + assert G1.args == G2.args + + X = MatrixSymbol('X', 2, 2) + G = GSE('U', 2) + h_pspace = RandomMatrixPSpace('P', model=density(G)) + H = RandomMatrixSymbol('H', 2, 2, pspace=h_pspace) + H2 = RandomMatrixSymbol('H', 2, 2, pspace=None) + assert H.doit() == H + + assert (2*H).xreplace({H: X}) == 2*X + assert (2*H).xreplace({H2: X}) == 2*H + assert (2*H2).xreplace({H: X}) == 2*H2 + assert (2*H2).xreplace({H2: X}) == 2*X diff --git a/lib/python3.12/site-packages/sympy/stats/tests/test_rv.py b/lib/python3.12/site-packages/sympy/stats/tests/test_rv.py new file mode 100644 index 0000000000000000000000000000000000000000..185756300556f2fe70b76c402113ec2bb2501ef4 --- /dev/null +++ b/lib/python3.12/site-packages/sympy/stats/tests/test_rv.py @@ -0,0 +1,441 @@ +from sympy.concrete.summations import Sum +from sympy.core.basic import Basic +from sympy.core.containers import Tuple +from sympy.core.function import Lambda +from sympy.core.numbers import (Rational, nan, oo, pi) +from sympy.core.relational import Eq +from sympy.core.singleton import S +from sympy.core.symbol import (Symbol, symbols) +from sympy.functions.combinatorial.factorials import (FallingFactorial, binomial) +from sympy.functions.elementary.exponential import (exp, log) +from sympy.functions.elementary.trigonometric import (cos, sin) +from sympy.functions.special.delta_functions import DiracDelta +from sympy.integrals.integrals import integrate +from sympy.logic.boolalg import (And, Or) +from sympy.matrices.dense import Matrix +from sympy.sets.sets import Interval +from sympy.tensor.indexed import Indexed +from sympy.stats import (Die, Normal, Exponential, FiniteRV, P, E, H, variance, + density, given, independent, dependent, where, pspace, GaussianUnitaryEnsemble, + random_symbols, sample, Geometric, factorial_moment, Binomial, Hypergeometric, + DiscreteUniform, Poisson, characteristic_function, moment_generating_function, + BernoulliProcess, Variance, Expectation, Probability, Covariance, covariance, cmoment, + moment, median) +from sympy.stats.rv import (IndependentProductPSpace, rs_swap, Density, NamedArgsMixin, + RandomSymbol, sample_iter, PSpace, is_random, RandomIndexedSymbol, RandomMatrixSymbol) +from sympy.testing.pytest import raises, skip, XFAIL, warns_deprecated_sympy +from sympy.external import import_module +from sympy.core.numbers import comp +from sympy.stats.frv_types import BernoulliDistribution +from sympy.core.symbol import Dummy +from sympy.functions.elementary.piecewise import Piecewise + +def test_where(): + X, Y = Die('X'), Die('Y') + Z = Normal('Z', 0, 1) + + assert where(Z**2 <= 1).set == Interval(-1, 1) + assert where(Z**2 <= 1).as_boolean() == Interval(-1, 1).as_relational(Z.symbol) + assert where(And(X > Y, Y > 4)).as_boolean() == And( + Eq(X.symbol, 6), Eq(Y.symbol, 5)) + + assert len(where(X < 3).set) == 2 + assert 1 in where(X < 3).set + + X, Y = Normal('X', 0, 1), Normal('Y', 0, 1) + assert where(And(X**2 <= 1, X >= 0)).set == Interval(0, 1) + XX = given(X, And(X**2 <= 1, X >= 0)) + assert XX.pspace.domain.set == Interval(0, 1) + assert XX.pspace.domain.as_boolean() == \ + And(0 <= X.symbol, X.symbol**2 <= 1, -oo < X.symbol, X.symbol < oo) + + with raises(TypeError): + XX = given(X, X + 3) + + +def test_random_symbols(): + X, Y = Normal('X', 0, 1), Normal('Y', 0, 1) + + assert set(random_symbols(2*X + 1)) == {X} + assert set(random_symbols(2*X + Y)) == {X, Y} + assert set(random_symbols(2*X + Y.symbol)) == {X} + assert set(random_symbols(2)) == set() + + +def test_characteristic_function(): + # Imports I from sympy + from sympy.core.numbers import I + X = Normal('X',0,1) + Y = DiscreteUniform('Y', [1,2,7]) + Z = Poisson('Z', 2) + t = symbols('_t') + P = Lambda(t, exp(-t**2/2)) + Q = Lambda(t, exp(7*t*I)/3 + exp(2*t*I)/3 + exp(t*I)/3) + R = Lambda(t, exp(2 * exp(t*I) - 2)) + + + assert characteristic_function(X).dummy_eq(P) + assert characteristic_function(Y).dummy_eq(Q) + assert characteristic_function(Z).dummy_eq(R) + + +def test_moment_generating_function(): + + X = Normal('X',0,1) + Y = DiscreteUniform('Y', [1,2,7]) + Z = Poisson('Z', 2) + t = symbols('_t') + P = Lambda(t, exp(t**2/2)) + Q = Lambda(t, (exp(7*t)/3 + exp(2*t)/3 + exp(t)/3)) + R = Lambda(t, exp(2 * exp(t) - 2)) + + + assert moment_generating_function(X).dummy_eq(P) + assert moment_generating_function(Y).dummy_eq(Q) + assert moment_generating_function(Z).dummy_eq(R) + +def test_sample_iter(): + + X = Normal('X',0,1) + Y = DiscreteUniform('Y', [1, 2, 7]) + Z = Poisson('Z', 2) + + scipy = import_module('scipy') + if not scipy: + skip('Scipy is not installed. Abort tests') + expr = X**2 + 3 + iterator = sample_iter(expr) + + expr2 = Y**2 + 5*Y + 4 + iterator2 = sample_iter(expr2) + + expr3 = Z**3 + 4 + iterator3 = sample_iter(expr3) + + def is_iterator(obj): + if ( + hasattr(obj, '__iter__') and + (hasattr(obj, 'next') or + hasattr(obj, '__next__')) and + callable(obj.__iter__) and + obj.__iter__() is obj + ): + return True + else: + return False + assert is_iterator(iterator) + assert is_iterator(iterator2) + assert is_iterator(iterator3) + +def test_pspace(): + X, Y = Normal('X', 0, 1), Normal('Y', 0, 1) + x = Symbol('x') + + raises(ValueError, lambda: pspace(5 + 3)) + raises(ValueError, lambda: pspace(x < 1)) + assert pspace(X) == X.pspace + assert pspace(2*X + 1) == X.pspace + assert pspace(2*X + Y) == IndependentProductPSpace(Y.pspace, X.pspace) + +def test_rs_swap(): + X = Normal('x', 0, 1) + Y = Exponential('y', 1) + + XX = Normal('x', 0, 2) + YY = Normal('y', 0, 3) + + expr = 2*X + Y + assert expr.subs(rs_swap((X, Y), (YY, XX))) == 2*XX + YY + + +def test_RandomSymbol(): + + X = Normal('x', 0, 1) + Y = Normal('x', 0, 2) + assert X.symbol == Y.symbol + assert X != Y + + assert X.name == X.symbol.name + + X = Normal('lambda', 0, 1) # make sure we can use protected terms + X = Normal('Lambda', 0, 1) # make sure we can use SymPy terms + + +def test_RandomSymbol_diff(): + X = Normal('x', 0, 1) + assert (2*X).diff(X) + + +def test_random_symbol_no_pspace(): + x = RandomSymbol(Symbol('x')) + assert x.pspace == PSpace() + +def test_overlap(): + X = Normal('x', 0, 1) + Y = Normal('x', 0, 2) + + raises(ValueError, lambda: P(X > Y)) + + +def test_IndependentProductPSpace(): + X = Normal('X', 0, 1) + Y = Normal('Y', 0, 1) + px = X.pspace + py = Y.pspace + assert pspace(X + Y) == IndependentProductPSpace(px, py) + assert pspace(X + Y) == IndependentProductPSpace(py, px) + + +def test_E(): + assert E(5) == 5 + + +def test_H(): + X = Normal('X', 0, 1) + D = Die('D', sides = 4) + G = Geometric('G', 0.5) + assert H(X, X > 0) == -log(2)/2 + S.Half + log(pi)/2 + assert H(D, D > 2) == log(2) + assert comp(H(G).evalf().round(2), 1.39) + + +def test_Sample(): + X = Die('X', 6) + Y = Normal('Y', 0, 1) + z = Symbol('z', integer=True) + + scipy = import_module('scipy') + if not scipy: + skip('Scipy is not installed. Abort tests') + assert sample(X) in [1, 2, 3, 4, 5, 6] + assert isinstance(sample(X + Y), float) + + assert P(X + Y > 0, Y < 0, numsamples=10).is_number + assert E(X + Y, numsamples=10).is_number + assert E(X**2 + Y, numsamples=10).is_number + assert E((X + Y)**2, numsamples=10).is_number + assert variance(X + Y, numsamples=10).is_number + + raises(TypeError, lambda: P(Y > z, numsamples=5)) + + assert P(sin(Y) <= 1, numsamples=10) == 1.0 + assert P(sin(Y) <= 1, cos(Y) < 1, numsamples=10) == 1.0 + + assert all(i in range(1, 7) for i in density(X, numsamples=10)) + assert all(i in range(4, 7) for i in density(X, X>3, numsamples=10)) + + numpy = import_module('numpy') + if not numpy: + skip('Numpy is not installed. Abort tests') + #Test Issue #21563: Output of sample must be a float or array + assert isinstance(sample(X), (numpy.int32, numpy.int64)) + assert isinstance(sample(Y), numpy.float64) + assert isinstance(sample(X, size=2), numpy.ndarray) + + with warns_deprecated_sympy(): + sample(X, numsamples=2) + +@XFAIL +def test_samplingE(): + scipy = import_module('scipy') + if not scipy: + skip('Scipy is not installed. Abort tests') + Y = Normal('Y', 0, 1) + z = Symbol('z', integer=True) + assert E(Sum(1/z**Y, (z, 1, oo)), Y > 2, numsamples=3).is_number + + +def test_given(): + X = Normal('X', 0, 1) + Y = Normal('Y', 0, 1) + A = given(X, True) + B = given(X, Y > 2) + + assert X == A == B + + +def test_factorial_moment(): + X = Poisson('X', 2) + Y = Binomial('Y', 2, S.Half) + Z = Hypergeometric('Z', 4, 2, 2) + assert factorial_moment(X, 2) == 4 + assert factorial_moment(Y, 2) == S.Half + assert factorial_moment(Z, 2) == Rational(1, 3) + + x, y, z, l = symbols('x y z l') + Y = Binomial('Y', 2, y) + Z = Hypergeometric('Z', 10, 2, 3) + assert factorial_moment(Y, l) == y**2*FallingFactorial( + 2, l) + 2*y*(1 - y)*FallingFactorial(1, l) + (1 - y)**2*\ + FallingFactorial(0, l) + assert factorial_moment(Z, l) == 7*FallingFactorial(0, l)/\ + 15 + 7*FallingFactorial(1, l)/15 + FallingFactorial(2, l)/15 + + +def test_dependence(): + X, Y = Die('X'), Die('Y') + assert independent(X, 2*Y) + assert not dependent(X, 2*Y) + + X, Y = Normal('X', 0, 1), Normal('Y', 0, 1) + assert independent(X, Y) + assert dependent(X, 2*X) + + # Create a dependency + XX, YY = given(Tuple(X, Y), Eq(X + Y, 3)) + assert dependent(XX, YY) + +def test_dependent_finite(): + X, Y = Die('X'), Die('Y') + # Dependence testing requires symbolic conditions which currently break + # finite random variables + assert dependent(X, Y + X) + + XX, YY = given(Tuple(X, Y), X + Y > 5) # Create a dependency + assert dependent(XX, YY) + + +def test_normality(): + X, Y = Normal('X', 0, 1), Normal('Y', 0, 1) + x = Symbol('x', real=True) + z = Symbol('z', real=True) + dens = density(X - Y, Eq(X + Y, z)) + + assert integrate(dens(x), (x, -oo, oo)) == 1 + + +def test_Density(): + X = Die('X', 6) + d = Density(X) + assert d.doit() == density(X) + +def test_NamedArgsMixin(): + class Foo(Basic, NamedArgsMixin): + _argnames = 'foo', 'bar' + + a = Foo(S(1), S(2)) + + assert a.foo == 1 + assert a.bar == 2 + + raises(AttributeError, lambda: a.baz) + + class Bar(Basic, NamedArgsMixin): + pass + + raises(AttributeError, lambda: Bar(S(1), S(2)).foo) + +def test_density_constant(): + assert density(3)(2) == 0 + assert density(3)(3) == DiracDelta(0) + +def test_cmoment_constant(): + assert variance(3) == 0 + assert cmoment(3, 3) == 0 + assert cmoment(3, 4) == 0 + x = Symbol('x') + assert variance(x) == 0 + assert cmoment(x, 15) == 0 + assert cmoment(x, 0) == 1 + +def test_moment_constant(): + assert moment(3, 0) == 1 + assert moment(3, 1) == 3 + assert moment(3, 2) == 9 + x = Symbol('x') + assert moment(x, 2) == x**2 + +def test_median_constant(): + assert median(3) == 3 + x = Symbol('x') + assert median(x) == x + +def test_real(): + x = Normal('x', 0, 1) + assert x.is_real + + +def test_issue_10052(): + X = Exponential('X', 3) + assert P(X < oo) == 1 + assert P(X > oo) == 0 + assert P(X < 2, X > oo) == 0 + assert P(X < oo, X > oo) == 0 + assert P(X < oo, X > 2) == 1 + assert P(X < 3, X == 2) == 0 + raises(ValueError, lambda: P(1)) + raises(ValueError, lambda: P(X < 1, 2)) + +def test_issue_11934(): + density = {0: .5, 1: .5} + X = FiniteRV('X', density) + assert E(X) == 0.5 + assert P( X>= 2) == 0 + +def test_issue_8129(): + X = Exponential('X', 4) + assert P(X >= X) == 1 + assert P(X > X) == 0 + assert P(X > X+1) == 0 + +def test_issue_12237(): + X = Normal('X', 0, 1) + Y = Normal('Y', 0, 1) + U = P(X > 0, X) + V = P(Y < 0, X) + W = P(X + Y > 0, X) + assert W == P(X + Y > 0, X) + assert U == BernoulliDistribution(S.Half, S.Zero, S.One) + assert V == S.Half + +def test_is_random(): + X = Normal('X', 0, 1) + Y = Normal('Y', 0, 1) + a, b = symbols('a, b') + G = GaussianUnitaryEnsemble('U', 2) + B = BernoulliProcess('B', 0.9) + assert not is_random(a) + assert not is_random(a + b) + assert not is_random(a * b) + assert not is_random(Matrix([a**2, b**2])) + assert is_random(X) + assert is_random(X**2 + Y) + assert is_random(Y + b**2) + assert is_random(Y > 5) + assert is_random(B[3] < 1) + assert is_random(G) + assert is_random(X * Y * B[1]) + assert is_random(Matrix([[X, B[2]], [G, Y]])) + assert is_random(Eq(X, 4)) + +def test_issue_12283(): + x = symbols('x') + X = RandomSymbol(x) + Y = RandomSymbol('Y') + Z = RandomMatrixSymbol('Z', 2, 1) + W = RandomMatrixSymbol('W', 2, 1) + RI = RandomIndexedSymbol(Indexed('RI', 3)) + assert pspace(Z) == PSpace() + assert pspace(RI) == PSpace() + assert pspace(X) == PSpace() + assert E(X) == Expectation(X) + assert P(Y > 3) == Probability(Y > 3) + assert variance(X) == Variance(X) + assert variance(RI) == Variance(RI) + assert covariance(X, Y) == Covariance(X, Y) + assert covariance(W, Z) == Covariance(W, Z) + +def test_issue_6810(): + X = Die('X', 6) + Y = Normal('Y', 0, 1) + assert P(Eq(X, 2)) == S(1)/6 + assert P(Eq(Y, 0)) == 0 + assert P(Or(X > 2, X < 3)) == 1 + assert P(And(X > 3, X > 2)) == S(1)/2 + +def test_issue_20286(): + n, p = symbols('n p') + B = Binomial('B', n, p) + k = Dummy('k', integer = True) + eq = Sum(Piecewise((-p**k*(1 - p)**(-k + n)*log(p**k*(1 - p)**(-k + n)*binomial(n, k))*binomial(n, k), (k >= 0) & (k <= n)), (nan, True)), (k, 0, n)) + assert eq.dummy_eq(H(B)) diff --git a/lib/python3.12/site-packages/sympy/stats/tests/test_stochastic_process.py b/lib/python3.12/site-packages/sympy/stats/tests/test_stochastic_process.py new file mode 100644 index 0000000000000000000000000000000000000000..d3d373821de7e28daf75c11d9d29d7bc5bdae1da --- /dev/null +++ b/lib/python3.12/site-packages/sympy/stats/tests/test_stochastic_process.py @@ -0,0 +1,763 @@ +from sympy.concrete.summations import Sum +from sympy.core.containers import Tuple +from sympy.core.function import Lambda +from sympy.core.numbers import (Float, Rational, oo, pi) +from sympy.core.relational import (Eq, Ge, Gt, Le, Lt, Ne) +from sympy.core.singleton import S +from sympy.core.symbol import (Symbol, symbols) +from sympy.functions.combinatorial.factorials import factorial +from sympy.functions.elementary.exponential import exp +from sympy.functions.elementary.integers import ceiling +from sympy.functions.elementary.miscellaneous import sqrt +from sympy.functions.elementary.piecewise import Piecewise +from sympy.functions.special.error_functions import erf +from sympy.functions.special.gamma_functions import (gamma, lowergamma) +from sympy.logic.boolalg import (And, Not) +from sympy.matrices.dense import Matrix +from sympy.matrices.expressions.matexpr import MatrixSymbol +from sympy.matrices.immutable import ImmutableMatrix +from sympy.sets.contains import Contains +from sympy.sets.fancysets import Range +from sympy.sets.sets import (FiniteSet, Interval) +from sympy.stats import (DiscreteMarkovChain, P, TransitionMatrixOf, E, + StochasticStateSpaceOf, variance, ContinuousMarkovChain, + BernoulliProcess, PoissonProcess, WienerProcess, + GammaProcess, sample_stochastic_process) +from sympy.stats.joint_rv import JointDistribution +from sympy.stats.joint_rv_types import JointDistributionHandmade +from sympy.stats.rv import RandomIndexedSymbol +from sympy.stats.symbolic_probability import Probability, Expectation +from sympy.testing.pytest import (raises, skip, ignore_warnings, + warns_deprecated_sympy) +from sympy.external import import_module +from sympy.stats.frv_types import BernoulliDistribution +from sympy.stats.drv_types import PoissonDistribution +from sympy.stats.crv_types import NormalDistribution, GammaDistribution +from sympy.core.symbol import Str + + +def test_DiscreteMarkovChain(): + + # pass only the name + X = DiscreteMarkovChain("X") + assert isinstance(X.state_space, Range) + assert X.index_set == S.Naturals0 + assert isinstance(X.transition_probabilities, MatrixSymbol) + t = symbols('t', positive=True, integer=True) + assert isinstance(X[t], RandomIndexedSymbol) + assert E(X[0]) == Expectation(X[0]) + raises(TypeError, lambda: DiscreteMarkovChain(1)) + raises(NotImplementedError, lambda: X(t)) + raises(NotImplementedError, lambda: X.communication_classes()) + raises(NotImplementedError, lambda: X.canonical_form()) + raises(NotImplementedError, lambda: X.decompose()) + + nz = Symbol('n', integer=True) + TZ = MatrixSymbol('M', nz, nz) + SZ = Range(nz) + YZ = DiscreteMarkovChain('Y', SZ, TZ) + assert P(Eq(YZ[2], 1), Eq(YZ[1], 0)) == TZ[0, 1] + + raises(ValueError, lambda: sample_stochastic_process(t)) + raises(ValueError, lambda: next(sample_stochastic_process(X))) + # pass name and state_space + # any hashable object should be a valid state + # states should be valid as a tuple/set/list/Tuple/Range + sym, rainy, cloudy, sunny = symbols('a Rainy Cloudy Sunny', real=True) + state_spaces = [(1, 2, 3), [Str('Hello'), sym, DiscreteMarkovChain("Y", (1,2,3))], + Tuple(S(1), exp(sym), Str('World'), sympify=False), Range(-1, 5, 2), + [rainy, cloudy, sunny]] + chains = [DiscreteMarkovChain("Y", state_space) for state_space in state_spaces] + + for i, Y in enumerate(chains): + assert isinstance(Y.transition_probabilities, MatrixSymbol) + assert Y.state_space == state_spaces[i] or Y.state_space == FiniteSet(*state_spaces[i]) + assert Y.number_of_states == 3 + + with ignore_warnings(UserWarning): # TODO: Restore tests once warnings are removed + assert P(Eq(Y[2], 1), Eq(Y[0], 2), evaluate=False) == Probability(Eq(Y[2], 1), Eq(Y[0], 2)) + assert E(Y[0]) == Expectation(Y[0]) + + raises(ValueError, lambda: next(sample_stochastic_process(Y))) + + raises(TypeError, lambda: DiscreteMarkovChain("Y", {1: 1})) + Y = DiscreteMarkovChain("Y", Range(1, t, 2)) + assert Y.number_of_states == ceiling((t-1)/2) + + # pass name and transition_probabilities + chains = [DiscreteMarkovChain("Y", trans_probs=Matrix([[]])), + DiscreteMarkovChain("Y", trans_probs=Matrix([[0, 1], [1, 0]])), + DiscreteMarkovChain("Y", trans_probs=Matrix([[pi, 1-pi], [sym, 1-sym]]))] + for Z in chains: + assert Z.number_of_states == Z.transition_probabilities.shape[0] + assert isinstance(Z.transition_probabilities, ImmutableMatrix) + + # pass name, state_space and transition_probabilities + T = Matrix([[0.5, 0.2, 0.3],[0.2, 0.5, 0.3],[0.2, 0.3, 0.5]]) + TS = MatrixSymbol('T', 3, 3) + Y = DiscreteMarkovChain("Y", [0, 1, 2], T) + YS = DiscreteMarkovChain("Y", ['One', 'Two', 3], TS) + assert Y.joint_distribution(1, Y[2], 3) == JointDistribution(Y[1], Y[2], Y[3]) + raises(ValueError, lambda: Y.joint_distribution(Y[1].symbol, Y[2].symbol)) + assert P(Eq(Y[3], 2), Eq(Y[1], 1)).round(2) == Float(0.36, 2) + assert (P(Eq(YS[3], 2), Eq(YS[1], 1)) - + (TS[0, 2]*TS[1, 0] + TS[1, 1]*TS[1, 2] + TS[1, 2]*TS[2, 2])).simplify() == 0 + assert P(Eq(YS[1], 1), Eq(YS[2], 2)) == Probability(Eq(YS[1], 1)) + assert P(Eq(YS[3], 3), Eq(YS[1], 1)) == TS[0, 2]*TS[1, 0] + TS[1, 1]*TS[1, 2] + TS[1, 2]*TS[2, 2] + TO = Matrix([[0.25, 0.75, 0],[0, 0.25, 0.75],[0.75, 0, 0.25]]) + assert P(Eq(Y[3], 2), Eq(Y[1], 1) & TransitionMatrixOf(Y, TO)).round(3) == Float(0.375, 3) + with ignore_warnings(UserWarning): ### TODO: Restore tests once warnings are removed + assert E(Y[3], evaluate=False) == Expectation(Y[3]) + assert E(Y[3], Eq(Y[2], 1)).round(2) == Float(1.1, 3) + TSO = MatrixSymbol('T', 4, 4) + raises(ValueError, lambda: str(P(Eq(YS[3], 2), Eq(YS[1], 1) & TransitionMatrixOf(YS, TSO)))) + raises(TypeError, lambda: DiscreteMarkovChain("Z", [0, 1, 2], symbols('M'))) + raises(ValueError, lambda: DiscreteMarkovChain("Z", [0, 1, 2], MatrixSymbol('T', 3, 4))) + raises(ValueError, lambda: E(Y[3], Eq(Y[2], 6))) + raises(ValueError, lambda: E(Y[2], Eq(Y[3], 1))) + + + # extended tests for probability queries + TO1 = Matrix([[Rational(1, 4), Rational(3, 4), 0],[Rational(1, 3), Rational(1, 3), Rational(1, 3)],[0, Rational(1, 4), Rational(3, 4)]]) + assert P(And(Eq(Y[2], 1), Eq(Y[1], 1), Eq(Y[0], 0)), + Eq(Probability(Eq(Y[0], 0)), Rational(1, 4)) & TransitionMatrixOf(Y, TO1)) == Rational(1, 16) + assert P(And(Eq(Y[2], 1), Eq(Y[1], 1), Eq(Y[0], 0)), TransitionMatrixOf(Y, TO1)) == \ + Probability(Eq(Y[0], 0))/4 + assert P(Lt(X[1], 2) & Gt(X[1], 0), Eq(X[0], 2) & + StochasticStateSpaceOf(X, [0, 1, 2]) & TransitionMatrixOf(X, TO1)) == Rational(1, 4) + assert P(Lt(X[1], 2) & Gt(X[1], 0), Eq(X[0], 2) & + StochasticStateSpaceOf(X, [S(0), '0', 1]) & TransitionMatrixOf(X, TO1)) == Rational(1, 4) + assert P(Ne(X[1], 2) & Ne(X[1], 1), Eq(X[0], 2) & + StochasticStateSpaceOf(X, [0, 1, 2]) & TransitionMatrixOf(X, TO1)) is S.Zero + assert P(Ne(X[1], 2) & Ne(X[1], 1), Eq(X[0], 2) & + StochasticStateSpaceOf(X, [S(0), '0', 1]) & TransitionMatrixOf(X, TO1)) is S.Zero + assert P(And(Eq(Y[2], 1), Eq(Y[1], 1), Eq(Y[0], 0)), Eq(Y[1], 1)) == 0.1*Probability(Eq(Y[0], 0)) + + # testing properties of Markov chain + TO2 = Matrix([[S.One, 0, 0],[Rational(1, 3), Rational(1, 3), Rational(1, 3)],[0, Rational(1, 4), Rational(3, 4)]]) + TO3 = Matrix([[Rational(1, 4), Rational(3, 4), 0],[Rational(1, 3), Rational(1, 3), Rational(1, 3)], [0, Rational(1, 4), Rational(3, 4)]]) + Y2 = DiscreteMarkovChain('Y', trans_probs=TO2) + Y3 = DiscreteMarkovChain('Y', trans_probs=TO3) + assert Y3.fundamental_matrix() == ImmutableMatrix([[176, 81, -132], [36, 141, -52], [-44, -39, 208]])/125 + assert Y2.is_absorbing_chain() == True + assert Y3.is_absorbing_chain() == False + assert Y2.canonical_form() == ([0, 1, 2], TO2) + assert Y3.canonical_form() == ([0, 1, 2], TO3) + assert Y2.decompose() == ([0, 1, 2], TO2[0:1, 0:1], TO2[1:3, 0:1], TO2[1:3, 1:3]) + assert Y3.decompose() == ([0, 1, 2], TO3, Matrix(0, 3, []), Matrix(0, 0, [])) + TO4 = Matrix([[Rational(1, 5), Rational(2, 5), Rational(2, 5)], [Rational(1, 10), S.Half, Rational(2, 5)], [Rational(3, 5), Rational(3, 10), Rational(1, 10)]]) + Y4 = DiscreteMarkovChain('Y', trans_probs=TO4) + w = ImmutableMatrix([[Rational(11, 39), Rational(16, 39), Rational(4, 13)]]) + assert Y4.limiting_distribution == w + assert Y4.is_regular() == True + assert Y4.is_ergodic() == True + TS1 = MatrixSymbol('T', 3, 3) + Y5 = DiscreteMarkovChain('Y', trans_probs=TS1) + assert Y5.limiting_distribution(w, TO4).doit() == True + assert Y5.stationary_distribution(condition_set=True).subs(TS1, TO4).contains(w).doit() == S.true + TO6 = Matrix([[S.One, 0, 0, 0, 0],[S.Half, 0, S.Half, 0, 0],[0, S.Half, 0, S.Half, 0], [0, 0, S.Half, 0, S.Half], [0, 0, 0, 0, 1]]) + Y6 = DiscreteMarkovChain('Y', trans_probs=TO6) + assert Y6.fundamental_matrix() == ImmutableMatrix([[Rational(3, 2), S.One, S.Half], [S.One, S(2), S.One], [S.Half, S.One, Rational(3, 2)]]) + assert Y6.absorbing_probabilities() == ImmutableMatrix([[Rational(3, 4), Rational(1, 4)], [S.Half, S.Half], [Rational(1, 4), Rational(3, 4)]]) + with warns_deprecated_sympy(): + Y6.absorbing_probabilites() + TO7 = Matrix([[Rational(1, 2), Rational(1, 4), Rational(1, 4)], [Rational(1, 2), 0, Rational(1, 2)], [Rational(1, 4), Rational(1, 4), Rational(1, 2)]]) + Y7 = DiscreteMarkovChain('Y', trans_probs=TO7) + assert Y7.is_absorbing_chain() == False + assert Y7.fundamental_matrix() == ImmutableMatrix([[Rational(86, 75), Rational(1, 25), Rational(-14, 75)], + [Rational(2, 25), Rational(21, 25), Rational(2, 25)], + [Rational(-14, 75), Rational(1, 25), Rational(86, 75)]]) + + # test for zero-sized matrix functionality + X = DiscreteMarkovChain('X', trans_probs=Matrix([[]])) + assert X.number_of_states == 0 + assert X.stationary_distribution() == Matrix([[]]) + assert X.communication_classes() == [] + assert X.canonical_form() == ([], Matrix([[]])) + assert X.decompose() == ([], Matrix([[]]), Matrix([[]]), Matrix([[]])) + assert X.is_regular() == False + assert X.is_ergodic() == False + + # test communication_class + # see https://drive.google.com/drive/folders/1HbxLlwwn2b3U8Lj7eb_ASIUb5vYaNIjg?usp=sharing + # tutorial 2.pdf + TO7 = Matrix([[0, 5, 5, 0, 0], + [0, 0, 0, 10, 0], + [5, 0, 5, 0, 0], + [0, 10, 0, 0, 0], + [0, 3, 0, 3, 4]])/10 + Y7 = DiscreteMarkovChain('Y', trans_probs=TO7) + tuples = Y7.communication_classes() + classes, recurrence, periods = list(zip(*tuples)) + assert classes == ([1, 3], [0, 2], [4]) + assert recurrence == (True, False, False) + assert periods == (2, 1, 1) + + TO8 = Matrix([[0, 0, 0, 10, 0, 0], + [5, 0, 5, 0, 0, 0], + [0, 4, 0, 0, 0, 6], + [10, 0, 0, 0, 0, 0], + [0, 10, 0, 0, 0, 0], + [0, 0, 0, 5, 5, 0]])/10 + Y8 = DiscreteMarkovChain('Y', trans_probs=TO8) + tuples = Y8.communication_classes() + classes, recurrence, periods = list(zip(*tuples)) + assert classes == ([0, 3], [1, 2, 5, 4]) + assert recurrence == (True, False) + assert periods == (2, 2) + + TO9 = Matrix([[2, 0, 0, 3, 0, 0, 3, 2, 0, 0], + [0, 10, 0, 0, 0, 0, 0, 0, 0, 0], + [0, 2, 2, 0, 0, 0, 0, 0, 3, 3], + [0, 0, 0, 3, 0, 0, 6, 1, 0, 0], + [0, 0, 0, 0, 5, 5, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 10, 0, 0, 0, 0], + [4, 0, 0, 5, 0, 0, 1, 0, 0, 0], + [2, 0, 0, 4, 0, 0, 2, 2, 0, 0], + [3, 0, 1, 0, 0, 0, 0, 0, 4, 2], + [0, 0, 4, 0, 0, 0, 0, 0, 3, 3]])/10 + Y9 = DiscreteMarkovChain('Y', trans_probs=TO9) + tuples = Y9.communication_classes() + classes, recurrence, periods = list(zip(*tuples)) + assert classes == ([0, 3, 6, 7], [1], [2, 8, 9], [5], [4]) + assert recurrence == (True, True, False, True, False) + assert periods == (1, 1, 1, 1, 1) + + # test canonical form + # see https://web.archive.org/web/20201230182007/https://www.dartmouth.edu/~chance/teaching_aids/books_articles/probability_book/Chapter11.pdf + # example 11.13 + T = Matrix([[1, 0, 0, 0, 0], + [S(1) / 2, 0, S(1) / 2, 0, 0], + [0, S(1) / 2, 0, S(1) / 2, 0], + [0, 0, S(1) / 2, 0, S(1) / 2], + [0, 0, 0, 0, S(1)]]) + DW = DiscreteMarkovChain('DW', [0, 1, 2, 3, 4], T) + states, A, B, C = DW.decompose() + assert states == [0, 4, 1, 2, 3] + assert A == Matrix([[1, 0], [0, 1]]) + assert B == Matrix([[S(1)/2, 0], [0, 0], [0, S(1)/2]]) + assert C == Matrix([[0, S(1)/2, 0], [S(1)/2, 0, S(1)/2], [0, S(1)/2, 0]]) + states, new_matrix = DW.canonical_form() + assert states == [0, 4, 1, 2, 3] + assert new_matrix == Matrix([[1, 0, 0, 0, 0], + [0, 1, 0, 0, 0], + [S(1)/2, 0, 0, S(1)/2, 0], + [0, 0, S(1)/2, 0, S(1)/2], + [0, S(1)/2, 0, S(1)/2, 0]]) + + # test regular and ergodic + # https://web.archive.org/web/20201230182007/https://www.dartmouth.edu/~chance/teaching_aids/books_articles/probability_book/Chapter11.pdf + T = Matrix([[0, 4, 0, 0, 0], + [1, 0, 3, 0, 0], + [0, 2, 0, 2, 0], + [0, 0, 3, 0, 1], + [0, 0, 0, 4, 0]])/4 + X = DiscreteMarkovChain('X', trans_probs=T) + assert not X.is_regular() + assert X.is_ergodic() + T = Matrix([[0, 1], [1, 0]]) + X = DiscreteMarkovChain('X', trans_probs=T) + assert not X.is_regular() + assert X.is_ergodic() + # http://www.math.wisc.edu/~valko/courses/331/MC2.pdf + T = Matrix([[2, 1, 1], + [2, 0, 2], + [1, 1, 2]])/4 + X = DiscreteMarkovChain('X', trans_probs=T) + assert X.is_regular() + assert X.is_ergodic() + # https://docs.ufpr.br/~lucambio/CE222/1S2014/Kemeny-Snell1976.pdf + T = Matrix([[1, 1], [1, 1]])/2 + X = DiscreteMarkovChain('X', trans_probs=T) + assert X.is_regular() + assert X.is_ergodic() + + # test is_absorbing_chain + T = Matrix([[0, 1, 0], + [1, 0, 0], + [0, 0, 1]]) + X = DiscreteMarkovChain('X', trans_probs=T) + assert not X.is_absorbing_chain() + # https://en.wikipedia.org/wiki/Absorbing_Markov_chain + T = Matrix([[1, 1, 0, 0], + [0, 1, 1, 0], + [1, 0, 0, 1], + [0, 0, 0, 2]])/2 + X = DiscreteMarkovChain('X', trans_probs=T) + assert X.is_absorbing_chain() + T = Matrix([[2, 0, 0, 0, 0], + [1, 0, 1, 0, 0], + [0, 1, 0, 1, 0], + [0, 0, 1, 0, 1], + [0, 0, 0, 0, 2]])/2 + X = DiscreteMarkovChain('X', trans_probs=T) + assert X.is_absorbing_chain() + + # test custom state space + Y10 = DiscreteMarkovChain('Y', [1, 2, 3], TO2) + tuples = Y10.communication_classes() + classes, recurrence, periods = list(zip(*tuples)) + assert classes == ([1], [2, 3]) + assert recurrence == (True, False) + assert periods == (1, 1) + assert Y10.canonical_form() == ([1, 2, 3], TO2) + assert Y10.decompose() == ([1, 2, 3], TO2[0:1, 0:1], TO2[1:3, 0:1], TO2[1:3, 1:3]) + + # testing miscellaneous queries + T = Matrix([[S.Half, Rational(1, 4), Rational(1, 4)], + [Rational(1, 3), 0, Rational(2, 3)], + [S.Half, S.Half, 0]]) + X = DiscreteMarkovChain('X', [0, 1, 2], T) + assert P(Eq(X[1], 2) & Eq(X[2], 1) & Eq(X[3], 0), + Eq(P(Eq(X[1], 0)), Rational(1, 4)) & Eq(P(Eq(X[1], 1)), Rational(1, 4))) == Rational(1, 12) + assert P(Eq(X[2], 1) | Eq(X[2], 2), Eq(X[1], 1)) == Rational(2, 3) + assert P(Eq(X[2], 1) & Eq(X[2], 2), Eq(X[1], 1)) is S.Zero + assert P(Ne(X[2], 2), Eq(X[1], 1)) == Rational(1, 3) + assert E(X[1]**2, Eq(X[0], 1)) == Rational(8, 3) + assert variance(X[1], Eq(X[0], 1)) == Rational(8, 9) + raises(ValueError, lambda: E(X[1], Eq(X[2], 1))) + raises(ValueError, lambda: DiscreteMarkovChain('X', [0, 1], T)) + + # testing miscellaneous queries with different state space + X = DiscreteMarkovChain('X', ['A', 'B', 'C'], T) + assert P(Eq(X[1], 2) & Eq(X[2], 1) & Eq(X[3], 0), + Eq(P(Eq(X[1], 0)), Rational(1, 4)) & Eq(P(Eq(X[1], 1)), Rational(1, 4))) == Rational(1, 12) + assert P(Eq(X[2], 1) | Eq(X[2], 2), Eq(X[1], 1)) == Rational(2, 3) + assert P(Eq(X[2], 1) & Eq(X[2], 2), Eq(X[1], 1)) is S.Zero + assert P(Ne(X[2], 2), Eq(X[1], 1)) == Rational(1, 3) + a = X.state_space.args[0] + c = X.state_space.args[2] + assert (E(X[1] ** 2, Eq(X[0], 1)) - (a**2/3 + 2*c**2/3)).simplify() == 0 + assert (variance(X[1], Eq(X[0], 1)) - (2*(-a/3 + c/3)**2/3 + (2*a/3 - 2*c/3)**2/3)).simplify() == 0 + raises(ValueError, lambda: E(X[1], Eq(X[2], 1))) + + #testing queries with multiple RandomIndexedSymbols + T = Matrix([[Rational(5, 10), Rational(3, 10), Rational(2, 10)], [Rational(2, 10), Rational(7, 10), Rational(1, 10)], [Rational(3, 10), Rational(3, 10), Rational(4, 10)]]) + Y = DiscreteMarkovChain("Y", [0, 1, 2], T) + assert P(Eq(Y[7], Y[5]), Eq(Y[2], 0)).round(5) == Float(0.44428, 5) + assert P(Gt(Y[3], Y[1]), Eq(Y[0], 0)).round(2) == Float(0.36, 2) + assert P(Le(Y[5], Y[10]), Eq(Y[4], 2)).round(6) == Float(0.583120, 6) + assert Float(P(Eq(Y[10], Y[5]), Eq(Y[4], 1)), 14) == Float(1 - P(Ne(Y[10], Y[5]), Eq(Y[4], 1)), 14) + assert Float(P(Gt(Y[8], Y[9]), Eq(Y[3], 2)), 14) == Float(1 - P(Le(Y[8], Y[9]), Eq(Y[3], 2)), 14) + assert Float(P(Lt(Y[1], Y[4]), Eq(Y[0], 0)), 14) == Float(1 - P(Ge(Y[1], Y[4]), Eq(Y[0], 0)), 14) + assert P(Eq(Y[5], Y[10]), Eq(Y[2], 1)) == P(Eq(Y[10], Y[5]), Eq(Y[2], 1)) + assert P(Gt(Y[1], Y[2]), Eq(Y[0], 1)) == P(Lt(Y[2], Y[1]), Eq(Y[0], 1)) + assert P(Ge(Y[7], Y[6]), Eq(Y[4], 1)) == P(Le(Y[6], Y[7]), Eq(Y[4], 1)) + + #test symbolic queries + a, b, c, d = symbols('a b c d') + T = Matrix([[Rational(1, 10), Rational(4, 10), Rational(5, 10)], [Rational(3, 10), Rational(4, 10), Rational(3, 10)], [Rational(7, 10), Rational(2, 10), Rational(1, 10)]]) + Y = DiscreteMarkovChain("Y", [0, 1, 2], T) + query = P(Eq(Y[a], b), Eq(Y[c], d)) + assert query.subs({a:10, b:2, c:5, d:1}).evalf().round(4) == P(Eq(Y[10], 2), Eq(Y[5], 1)).round(4) + assert query.subs({a:15, b:0, c:10, d:1}).evalf().round(4) == P(Eq(Y[15], 0), Eq(Y[10], 1)).round(4) + query_gt = P(Gt(Y[a], b), Eq(Y[c], d)) + query_le = P(Le(Y[a], b), Eq(Y[c], d)) + assert query_gt.subs({a:5, b:2, c:1, d:0}).evalf() + query_le.subs({a:5, b:2, c:1, d:0}).evalf() == 1.0 + query_ge = P(Ge(Y[a], b), Eq(Y[c], d)) + query_lt = P(Lt(Y[a], b), Eq(Y[c], d)) + assert query_ge.subs({a:4, b:1, c:0, d:2}).evalf() + query_lt.subs({a:4, b:1, c:0, d:2}).evalf() == 1.0 + + #test issue 20078 + assert (2*Y[1] + 3*Y[1]).simplify() == 5*Y[1] + assert (2*Y[1] - 3*Y[1]).simplify() == -Y[1] + assert (2*(0.25*Y[1])).simplify() == 0.5*Y[1] + assert ((2*Y[1]) * (0.25*Y[1])).simplify() == 0.5*Y[1]**2 + assert (Y[1]**2 + Y[1]**3).simplify() == (Y[1] + 1)*Y[1]**2 + +def test_sample_stochastic_process(): + if not import_module('scipy'): + skip('SciPy Not installed. Skip sampling tests') + import random + random.seed(0) + numpy = import_module('numpy') + if numpy: + numpy.random.seed(0) # scipy uses numpy to sample so to set its seed + T = Matrix([[0.5, 0.2, 0.3],[0.2, 0.5, 0.3],[0.2, 0.3, 0.5]]) + Y = DiscreteMarkovChain("Y", [0, 1, 2], T) + for samps in range(10): + assert next(sample_stochastic_process(Y)) in Y.state_space + Z = DiscreteMarkovChain("Z", ['1', 1, 0], T) + for samps in range(10): + assert next(sample_stochastic_process(Z)) in Z.state_space + + T = Matrix([[S.Half, Rational(1, 4), Rational(1, 4)], + [Rational(1, 3), 0, Rational(2, 3)], + [S.Half, S.Half, 0]]) + X = DiscreteMarkovChain('X', [0, 1, 2], T) + for samps in range(10): + assert next(sample_stochastic_process(X)) in X.state_space + W = DiscreteMarkovChain('W', [1, pi, oo], T) + for samps in range(10): + assert next(sample_stochastic_process(W)) in W.state_space + + +def test_ContinuousMarkovChain(): + T1 = Matrix([[S(-2), S(2), S.Zero], + [S.Zero, S.NegativeOne, S.One], + [Rational(3, 2), Rational(3, 2), S(-3)]]) + C1 = ContinuousMarkovChain('C', [0, 1, 2], T1) + assert C1.limiting_distribution() == ImmutableMatrix([[Rational(3, 19), Rational(12, 19), Rational(4, 19)]]) + + T2 = Matrix([[-S.One, S.One, S.Zero], [S.One, -S.One, S.Zero], [S.Zero, S.One, -S.One]]) + C2 = ContinuousMarkovChain('C', [0, 1, 2], T2) + A, t = C2.generator_matrix, symbols('t', positive=True) + assert C2.transition_probabilities(A)(t) == Matrix([[S.Half + exp(-2*t)/2, S.Half - exp(-2*t)/2, 0], + [S.Half - exp(-2*t)/2, S.Half + exp(-2*t)/2, 0], + [S.Half - exp(-t) + exp(-2*t)/2, S.Half - exp(-2*t)/2, exp(-t)]]) + with ignore_warnings(UserWarning): ### TODO: Restore tests once warnings are removed + assert P(Eq(C2(1), 1), Eq(C2(0), 1), evaluate=False) == Probability(Eq(C2(1), 1), Eq(C2(0), 1)) + assert P(Eq(C2(1), 1), Eq(C2(0), 1)) == exp(-2)/2 + S.Half + assert P(Eq(C2(1), 0) & Eq(C2(2), 1) & Eq(C2(3), 1), + Eq(P(Eq(C2(1), 0)), S.Half)) == (Rational(1, 4) - exp(-2)/4)*(exp(-2)/2 + S.Half) + assert P(Not(Eq(C2(1), 0) & Eq(C2(2), 1) & Eq(C2(3), 2)) | + (Eq(C2(1), 0) & Eq(C2(2), 1) & Eq(C2(3), 2)), + Eq(P(Eq(C2(1), 0)), Rational(1, 4)) & Eq(P(Eq(C2(1), 1)), Rational(1, 4))) is S.One + assert E(C2(Rational(3, 2)), Eq(C2(0), 2)) == -exp(-3)/2 + 2*exp(Rational(-3, 2)) + S.Half + assert variance(C2(Rational(3, 2)), Eq(C2(0), 1)) == ((S.Half - exp(-3)/2)**2*(exp(-3)/2 + S.Half) + + (Rational(-1, 2) - exp(-3)/2)**2*(S.Half - exp(-3)/2)) + raises(KeyError, lambda: P(Eq(C2(1), 0), Eq(P(Eq(C2(1), 1)), S.Half))) + assert P(Eq(C2(1), 0), Eq(P(Eq(C2(5), 1)), S.Half)) == Probability(Eq(C2(1), 0)) + TS1 = MatrixSymbol('G', 3, 3) + CS1 = ContinuousMarkovChain('C', [0, 1, 2], TS1) + A = CS1.generator_matrix + assert CS1.transition_probabilities(A)(t) == exp(t*A) + + C3 = ContinuousMarkovChain('C', [Symbol('0'), Symbol('1'), Symbol('2')], T2) + assert P(Eq(C3(1), 1), Eq(C3(0), 1)) == exp(-2)/2 + S.Half + assert P(Eq(C3(1), Symbol('1')), Eq(C3(0), Symbol('1'))) == exp(-2)/2 + S.Half + + #test probability queries + G = Matrix([[-S(1), Rational(1, 10), Rational(9, 10)], [Rational(2, 5), -S(1), Rational(3, 5)], [Rational(1, 2), Rational(1, 2), -S(1)]]) + C = ContinuousMarkovChain('C', state_space=[0, 1, 2], gen_mat=G) + assert P(Eq(C(7.385), C(3.19)), Eq(C(0.862), 0)).round(5) == Float(0.35469, 5) + assert P(Gt(C(98.715), C(19.807)), Eq(C(11.314), 2)).round(5) == Float(0.32452, 5) + assert P(Le(C(5.9), C(10.112)), Eq(C(4), 1)).round(6) == Float(0.675214, 6) + assert Float(P(Eq(C(7.32), C(2.91)), Eq(C(2.63), 1)), 14) == Float(1 - P(Ne(C(7.32), C(2.91)), Eq(C(2.63), 1)), 14) + assert Float(P(Gt(C(3.36), C(1.101)), Eq(C(0.8), 2)), 14) == Float(1 - P(Le(C(3.36), C(1.101)), Eq(C(0.8), 2)), 14) + assert Float(P(Lt(C(4.9), C(2.79)), Eq(C(1.61), 0)), 14) == Float(1 - P(Ge(C(4.9), C(2.79)), Eq(C(1.61), 0)), 14) + assert P(Eq(C(5.243), C(10.912)), Eq(C(2.174), 1)) == P(Eq(C(10.912), C(5.243)), Eq(C(2.174), 1)) + assert P(Gt(C(2.344), C(9.9)), Eq(C(1.102), 1)) == P(Lt(C(9.9), C(2.344)), Eq(C(1.102), 1)) + assert P(Ge(C(7.87), C(1.008)), Eq(C(0.153), 1)) == P(Le(C(1.008), C(7.87)), Eq(C(0.153), 1)) + + #test symbolic queries + a, b, c, d = symbols('a b c d') + query = P(Eq(C(a), b), Eq(C(c), d)) + assert query.subs({a:3.65, b:2, c:1.78, d:1}).evalf().round(10) == P(Eq(C(3.65), 2), Eq(C(1.78), 1)).round(10) + query_gt = P(Gt(C(a), b), Eq(C(c), d)) + query_le = P(Le(C(a), b), Eq(C(c), d)) + assert query_gt.subs({a:13.2, b:0, c:3.29, d:2}).evalf() + query_le.subs({a:13.2, b:0, c:3.29, d:2}).evalf() == 1.0 + query_ge = P(Ge(C(a), b), Eq(C(c), d)) + query_lt = P(Lt(C(a), b), Eq(C(c), d)) + assert query_ge.subs({a:7.43, b:1, c:1.45, d:0}).evalf() + query_lt.subs({a:7.43, b:1, c:1.45, d:0}).evalf() == 1.0 + + #test issue 20078 + assert (2*C(1) + 3*C(1)).simplify() == 5*C(1) + assert (2*C(1) - 3*C(1)).simplify() == -C(1) + assert (2*(0.25*C(1))).simplify() == 0.5*C(1) + assert (2*C(1) * 0.25*C(1)).simplify() == 0.5*C(1)**2 + assert (C(1)**2 + C(1)**3).simplify() == (C(1) + 1)*C(1)**2 + +def test_BernoulliProcess(): + + B = BernoulliProcess("B", p=0.6, success=1, failure=0) + assert B.state_space == FiniteSet(0, 1) + assert B.index_set == S.Naturals0 + assert B.success == 1 + assert B.failure == 0 + + X = BernoulliProcess("X", p=Rational(1,3), success='H', failure='T') + assert X.state_space == FiniteSet('H', 'T') + H, T = symbols("H,T") + assert E(X[1]+X[2]*X[3]) == H**2/9 + 4*H*T/9 + H/3 + 4*T**2/9 + 2*T/3 + + t, x = symbols('t, x', positive=True, integer=True) + assert isinstance(B[t], RandomIndexedSymbol) + + raises(ValueError, lambda: BernoulliProcess("X", p=1.1, success=1, failure=0)) + raises(NotImplementedError, lambda: B(t)) + + raises(IndexError, lambda: B[-3]) + assert B.joint_distribution(B[3], B[9]) == JointDistributionHandmade(Lambda((B[3], B[9]), + Piecewise((0.6, Eq(B[3], 1)), (0.4, Eq(B[3], 0)), (0, True)) + *Piecewise((0.6, Eq(B[9], 1)), (0.4, Eq(B[9], 0)), (0, True)))) + + assert B.joint_distribution(2, B[4]) == JointDistributionHandmade(Lambda((B[2], B[4]), + Piecewise((0.6, Eq(B[2], 1)), (0.4, Eq(B[2], 0)), (0, True)) + *Piecewise((0.6, Eq(B[4], 1)), (0.4, Eq(B[4], 0)), (0, True)))) + + # Test for the sum distribution of Bernoulli Process RVs + Y = B[1] + B[2] + B[3] + assert P(Eq(Y, 0)).round(2) == Float(0.06, 1) + assert P(Eq(Y, 2)).round(2) == Float(0.43, 2) + assert P(Eq(Y, 4)).round(2) == 0 + assert P(Gt(Y, 1)).round(2) == Float(0.65, 2) + # Test for independency of each Random Indexed variable + assert P(Eq(B[1], 0) & Eq(B[2], 1) & Eq(B[3], 0) & Eq(B[4], 1)).round(2) == Float(0.06, 1) + + assert E(2 * B[1] + B[2]).round(2) == Float(1.80, 3) + assert E(2 * B[1] + B[2] + 5).round(2) == Float(6.80, 3) + assert E(B[2] * B[4] + B[10]).round(2) == Float(0.96, 2) + assert E(B[2] > 0, Eq(B[1],1) & Eq(B[2],1)).round(2) == Float(0.60,2) + assert E(B[1]) == 0.6 + assert P(B[1] > 0).round(2) == Float(0.60, 2) + assert P(B[1] < 1).round(2) == Float(0.40, 2) + assert P(B[1] > 0, B[2] <= 1).round(2) == Float(0.60, 2) + assert P(B[12] * B[5] > 0).round(2) == Float(0.36, 2) + assert P(B[12] * B[5] > 0, B[4] < 1).round(2) == Float(0.36, 2) + assert P(Eq(B[2], 1), B[2] > 0) == 1.0 + assert P(Eq(B[5], 3)) == 0 + assert P(Eq(B[1], 1), B[1] < 0) == 0 + assert P(B[2] > 0, Eq(B[2], 1)) == 1 + assert P(B[2] < 0, Eq(B[2], 1)) == 0 + assert P(B[2] > 0, B[2]==7) == 0 + assert P(B[5] > 0, B[5]) == BernoulliDistribution(0.6, 0, 1) + raises(ValueError, lambda: P(3)) + raises(ValueError, lambda: P(B[3] > 0, 3)) + + # test issue 19456 + expr = Sum(B[t], (t, 0, 4)) + expr2 = Sum(B[t], (t, 1, 3)) + expr3 = Sum(B[t]**2, (t, 1, 3)) + assert expr.doit() == B[0] + B[1] + B[2] + B[3] + B[4] + assert expr2.doit() == Y + assert expr3.doit() == B[1]**2 + B[2]**2 + B[3]**2 + assert B[2*t].free_symbols == {B[2*t], t} + assert B[4].free_symbols == {B[4]} + assert B[x*t].free_symbols == {B[x*t], x, t} + + #test issue 20078 + assert (2*B[t] + 3*B[t]).simplify() == 5*B[t] + assert (2*B[t] - 3*B[t]).simplify() == -B[t] + assert (2*(0.25*B[t])).simplify() == 0.5*B[t] + assert (2*B[t] * 0.25*B[t]).simplify() == 0.5*B[t]**2 + assert (B[t]**2 + B[t]**3).simplify() == (B[t] + 1)*B[t]**2 + +def test_PoissonProcess(): + X = PoissonProcess("X", 3) + assert X.state_space == S.Naturals0 + assert X.index_set == Interval(0, oo) + assert X.lamda == 3 + + t, d, x, y = symbols('t d x y', positive=True) + assert isinstance(X(t), RandomIndexedSymbol) + assert X.distribution(t) == PoissonDistribution(3*t) + with warns_deprecated_sympy(): + X.distribution(X(t)) + raises(ValueError, lambda: PoissonProcess("X", -1)) + raises(NotImplementedError, lambda: X[t]) + raises(IndexError, lambda: X(-5)) + + assert X.joint_distribution(X(2), X(3)) == JointDistributionHandmade(Lambda((X(2), X(3)), + 6**X(2)*9**X(3)*exp(-15)/(factorial(X(2))*factorial(X(3))))) + + assert X.joint_distribution(4, 6) == JointDistributionHandmade(Lambda((X(4), X(6)), + 12**X(4)*18**X(6)*exp(-30)/(factorial(X(4))*factorial(X(6))))) + + assert P(X(t) < 1) == exp(-3*t) + assert P(Eq(X(t), 0), Contains(t, Interval.Lopen(3, 5))) == exp(-6) # exp(-2*lamda) + res = P(Eq(X(t), 1), Contains(t, Interval.Lopen(3, 4))) + assert res == 3*exp(-3) + + # Equivalent to P(Eq(X(t), 1))**4 because of non-overlapping intervals + assert P(Eq(X(t), 1) & Eq(X(d), 1) & Eq(X(x), 1) & Eq(X(y), 1), Contains(t, Interval.Lopen(0, 1)) + & Contains(d, Interval.Lopen(1, 2)) & Contains(x, Interval.Lopen(2, 3)) + & Contains(y, Interval.Lopen(3, 4))) == res**4 + + # Return Probability because of overlapping intervals + assert P(Eq(X(t), 2) & Eq(X(d), 3), Contains(t, Interval.Lopen(0, 2)) + & Contains(d, Interval.Ropen(2, 4))) == \ + Probability(Eq(X(d), 3) & Eq(X(t), 2), Contains(t, Interval.Lopen(0, 2)) + & Contains(d, Interval.Ropen(2, 4))) + + raises(ValueError, lambda: P(Eq(X(t), 2) & Eq(X(d), 3), + Contains(t, Interval.Lopen(0, 4)) & Contains(d, Interval.Lopen(3, oo)))) # no bound on d + assert P(Eq(X(3), 2)) == 81*exp(-9)/2 + assert P(Eq(X(t), 2), Contains(t, Interval.Lopen(0, 5))) == 225*exp(-15)/2 + + # Check that probability works correctly by adding it to 1 + res1 = P(X(t) <= 3, Contains(t, Interval.Lopen(0, 5))) + res2 = P(X(t) > 3, Contains(t, Interval.Lopen(0, 5))) + assert res1 == 691*exp(-15) + assert (res1 + res2).simplify() == 1 + + # Check Not and Or + assert P(Not(Eq(X(t), 2) & (X(d) > 3)), Contains(t, Interval.Ropen(2, 4)) & \ + Contains(d, Interval.Lopen(7, 8))).simplify() == -18*exp(-6) + 234*exp(-9) + 1 + assert P(Eq(X(t), 2) | Ne(X(t), 4), Contains(t, Interval.Ropen(2, 4))) == 1 - 36*exp(-6) + raises(ValueError, lambda: P(X(t) > 2, X(t) + X(d))) + assert E(X(t)) == 3*t # property of the distribution at a given timestamp + assert E(X(t)**2 + X(d)*2 + X(y)**3, Contains(t, Interval.Lopen(0, 1)) + & Contains(d, Interval.Lopen(1, 2)) & Contains(y, Interval.Ropen(3, 4))) == 75 + assert E(X(t)**2, Contains(t, Interval.Lopen(0, 1))) == 12 + assert E(x*(X(t) + X(d))*(X(t)**2+X(d)**2), Contains(t, Interval.Lopen(0, 1)) + & Contains(d, Interval.Ropen(1, 2))) == \ + Expectation(x*(X(d) + X(t))*(X(d)**2 + X(t)**2), Contains(t, Interval.Lopen(0, 1)) + & Contains(d, Interval.Ropen(1, 2))) + + # Value Error because of infinite time bound + raises(ValueError, lambda: E(X(t)**3, Contains(t, Interval.Lopen(1, oo)))) + + # Equivalent to E(X(t)**2) - E(X(d)**2) == E(X(1)**2) - E(X(1)**2) == 0 + assert E((X(t) + X(d))*(X(t) - X(d)), Contains(t, Interval.Lopen(0, 1)) + & Contains(d, Interval.Lopen(1, 2))) == 0 + assert E(X(2) + x*E(X(5))) == 15*x + 6 + assert E(x*X(1) + y) == 3*x + y + assert P(Eq(X(1), 2) & Eq(X(t), 3), Contains(t, Interval.Lopen(1, 2))) == 81*exp(-6)/4 + Y = PoissonProcess("Y", 6) + Z = X + Y + assert Z.lamda == X.lamda + Y.lamda == 9 + raises(ValueError, lambda: X + 5) # should be added be only PoissonProcess instance + N, M = Z.split(4, 5) + assert N.lamda == 4 + assert M.lamda == 5 + raises(ValueError, lambda: Z.split(3, 2)) # 2+3 != 9 + + raises(ValueError, lambda :P(Eq(X(t), 0), Contains(t, Interval.Lopen(1, 3)) & Eq(X(1), 0))) + # check if it handles queries with two random variables in one args + res1 = P(Eq(N(3), N(5))) + assert res1 == P(Eq(N(t), 0), Contains(t, Interval(3, 5))) + res2 = P(N(3) > N(1)) + assert res2 == P((N(t) > 0), Contains(t, Interval(1, 3))) + assert P(N(3) < N(1)) == 0 # condition is not possible + res3 = P(N(3) <= N(1)) # holds only for Eq(N(3), N(1)) + assert res3 == P(Eq(N(t), 0), Contains(t, Interval(1, 3))) + + # tests from https://www.probabilitycourse.com/chapter11/11_1_2_basic_concepts_of_the_poisson_process.php + X = PoissonProcess('X', 10) # 11.1 + assert P(Eq(X(S(1)/3), 3) & Eq(X(1), 10)) == exp(-10)*Rational(8000000000, 11160261) + assert P(Eq(X(1), 1), Eq(X(S(1)/3), 3)) == 0 + assert P(Eq(X(1), 10), Eq(X(S(1)/3), 3)) == P(Eq(X(S(2)/3), 7)) + + X = PoissonProcess('X', 2) # 11.2 + assert P(X(S(1)/2) < 1) == exp(-1) + assert P(X(3) < 1, Eq(X(1), 0)) == exp(-4) + assert P(Eq(X(4), 3), Eq(X(2), 3)) == exp(-4) + + X = PoissonProcess('X', 3) + assert P(Eq(X(2), 5) & Eq(X(1), 2)) == Rational(81, 4)*exp(-6) + + # check few properties + assert P(X(2) <= 3, X(1)>=1) == 3*P(Eq(X(1), 0)) + 2*P(Eq(X(1), 1)) + P(Eq(X(1), 2)) + assert P(X(2) <= 3, X(1) > 1) == 2*P(Eq(X(1), 0)) + 1*P(Eq(X(1), 1)) + assert P(Eq(X(2), 5) & Eq(X(1), 2)) == P(Eq(X(1), 3))*P(Eq(X(1), 2)) + assert P(Eq(X(3), 4), Eq(X(1), 3)) == P(Eq(X(2), 1)) + + #test issue 20078 + assert (2*X(t) + 3*X(t)).simplify() == 5*X(t) + assert (2*X(t) - 3*X(t)).simplify() == -X(t) + assert (2*(0.25*X(t))).simplify() == 0.5*X(t) + assert (2*X(t) * 0.25*X(t)).simplify() == 0.5*X(t)**2 + assert (X(t)**2 + X(t)**3).simplify() == (X(t) + 1)*X(t)**2 + +def test_WienerProcess(): + X = WienerProcess("X") + assert X.state_space == S.Reals + assert X.index_set == Interval(0, oo) + + t, d, x, y = symbols('t d x y', positive=True) + assert isinstance(X(t), RandomIndexedSymbol) + assert X.distribution(t) == NormalDistribution(0, sqrt(t)) + with warns_deprecated_sympy(): + X.distribution(X(t)) + raises(ValueError, lambda: PoissonProcess("X", -1)) + raises(NotImplementedError, lambda: X[t]) + raises(IndexError, lambda: X(-2)) + + assert X.joint_distribution(X(2), X(3)) == JointDistributionHandmade( + Lambda((X(2), X(3)), sqrt(6)*exp(-X(2)**2/4)*exp(-X(3)**2/6)/(12*pi))) + assert X.joint_distribution(4, 6) == JointDistributionHandmade( + Lambda((X(4), X(6)), sqrt(6)*exp(-X(4)**2/8)*exp(-X(6)**2/12)/(24*pi))) + + assert P(X(t) < 3).simplify() == erf(3*sqrt(2)/(2*sqrt(t)))/2 + S(1)/2 + assert P(X(t) > 2, Contains(t, Interval.Lopen(3, 7))).simplify() == S(1)/2 -\ + erf(sqrt(2)/2)/2 + + # Equivalent to P(X(1)>1)**4 + assert P((X(t) > 4) & (X(d) > 3) & (X(x) > 2) & (X(y) > 1), + Contains(t, Interval.Lopen(0, 1)) & Contains(d, Interval.Lopen(1, 2)) + & Contains(x, Interval.Lopen(2, 3)) & Contains(y, Interval.Lopen(3, 4))).simplify() ==\ + (1 - erf(sqrt(2)/2))*(1 - erf(sqrt(2)))*(1 - erf(3*sqrt(2)/2))*(1 - erf(2*sqrt(2)))/16 + + # Contains an overlapping interval so, return Probability + assert P((X(t)< 2) & (X(d)> 3), Contains(t, Interval.Lopen(0, 2)) + & Contains(d, Interval.Ropen(2, 4))) == Probability((X(d) > 3) & (X(t) < 2), + Contains(d, Interval.Ropen(2, 4)) & Contains(t, Interval.Lopen(0, 2))) + + assert str(P(Not((X(t) < 5) & (X(d) > 3)), Contains(t, Interval.Ropen(2, 4)) & + Contains(d, Interval.Lopen(7, 8))).simplify()) == \ + '-(1 - erf(3*sqrt(2)/2))*(2 - erfc(5/2))/4 + 1' + # Distribution has mean 0 at each timestamp + assert E(X(t)) == 0 + assert E(x*(X(t) + X(d))*(X(t)**2+X(d)**2), Contains(t, Interval.Lopen(0, 1)) + & Contains(d, Interval.Ropen(1, 2))) == Expectation(x*(X(d) + X(t))*(X(d)**2 + X(t)**2), + Contains(d, Interval.Ropen(1, 2)) & Contains(t, Interval.Lopen(0, 1))) + assert E(X(t) + x*E(X(3))) == 0 + + #test issue 20078 + assert (2*X(t) + 3*X(t)).simplify() == 5*X(t) + assert (2*X(t) - 3*X(t)).simplify() == -X(t) + assert (2*(0.25*X(t))).simplify() == 0.5*X(t) + assert (2*X(t) * 0.25*X(t)).simplify() == 0.5*X(t)**2 + assert (X(t)**2 + X(t)**3).simplify() == (X(t) + 1)*X(t)**2 + + +def test_GammaProcess_symbolic(): + t, d, x, y, g, l = symbols('t d x y g l', positive=True) + X = GammaProcess("X", l, g) + + raises(NotImplementedError, lambda: X[t]) + raises(IndexError, lambda: X(-1)) + assert isinstance(X(t), RandomIndexedSymbol) + assert X.state_space == Interval(0, oo) + assert X.distribution(t) == GammaDistribution(g*t, 1/l) + with warns_deprecated_sympy(): + X.distribution(X(t)) + assert X.joint_distribution(5, X(3)) == JointDistributionHandmade(Lambda( + (X(5), X(3)), l**(8*g)*exp(-l*X(3))*exp(-l*X(5))*X(3)**(3*g - 1)*X(5)**(5*g + - 1)/(gamma(3*g)*gamma(5*g)))) + # property of the gamma process at any given timestamp + assert E(X(t)) == g*t/l + assert variance(X(t)).simplify() == g*t/l**2 + + # Equivalent to E(2*X(1)) + E(X(1)**2) + E(X(1)**3), where E(X(1)) == g/l + assert E(X(t)**2 + X(d)*2 + X(y)**3, Contains(t, Interval.Lopen(0, 1)) + & Contains(d, Interval.Lopen(1, 2)) & Contains(y, Interval.Ropen(3, 4))) == \ + 2*g/l + (g**2 + g)/l**2 + (g**3 + 3*g**2 + 2*g)/l**3 + + assert P(X(t) > 3, Contains(t, Interval.Lopen(3, 4))).simplify() == \ + 1 - lowergamma(g, 3*l)/gamma(g) # equivalent to P(X(1)>3) + + + #test issue 20078 + assert (2*X(t) + 3*X(t)).simplify() == 5*X(t) + assert (2*X(t) - 3*X(t)).simplify() == -X(t) + assert (2*(0.25*X(t))).simplify() == 0.5*X(t) + assert (2*X(t) * 0.25*X(t)).simplify() == 0.5*X(t)**2 + assert (X(t)**2 + X(t)**3).simplify() == (X(t) + 1)*X(t)**2 +def test_GammaProcess_numeric(): + t, d, x, y = symbols('t d x y', positive=True) + X = GammaProcess("X", 1, 2) + assert X.state_space == Interval(0, oo) + assert X.index_set == Interval(0, oo) + assert X.lamda == 1 + assert X.gamma == 2 + + raises(ValueError, lambda: GammaProcess("X", -1, 2)) + raises(ValueError, lambda: GammaProcess("X", 0, -2)) + raises(ValueError, lambda: GammaProcess("X", -1, -2)) + + # all are independent because of non-overlapping intervals + assert P((X(t) > 4) & (X(d) > 3) & (X(x) > 2) & (X(y) > 1), Contains(t, + Interval.Lopen(0, 1)) & Contains(d, Interval.Lopen(1, 2)) & Contains(x, + Interval.Lopen(2, 3)) & Contains(y, Interval.Lopen(3, 4))).simplify() == \ + 120*exp(-10) + + # Check working with Not and Or + assert P(Not((X(t) < 5) & (X(d) > 3)), Contains(t, Interval.Ropen(2, 4)) & + Contains(d, Interval.Lopen(7, 8))).simplify() == -4*exp(-3) + 472*exp(-8)/3 + 1 + assert P((X(t) > 2) | (X(t) < 4), Contains(t, Interval.Ropen(1, 4))).simplify() == \ + -643*exp(-4)/15 + 109*exp(-2)/15 + 1 + + assert E(X(t)) == 2*t # E(X(t)) == gamma*t/l + assert E(X(2) + x*E(X(5))) == 10*x + 4 diff --git a/lib/python3.12/site-packages/sympy/stats/tests/test_symbolic_multivariate.py b/lib/python3.12/site-packages/sympy/stats/tests/test_symbolic_multivariate.py new file mode 100644 index 0000000000000000000000000000000000000000..79979e20a6f10d2a2ddfe85ce4c8df145e98c3fd --- /dev/null +++ b/lib/python3.12/site-packages/sympy/stats/tests/test_symbolic_multivariate.py @@ -0,0 +1,172 @@ +from sympy.stats import Expectation, Normal, Variance, Covariance +from sympy.testing.pytest import raises +from sympy.core.symbol import symbols +from sympy.matrices.exceptions import ShapeError +from sympy.matrices.dense import Matrix +from sympy.matrices.expressions.matexpr import MatrixSymbol +from sympy.matrices.expressions.special import ZeroMatrix +from sympy.stats.rv import RandomMatrixSymbol +from sympy.stats.symbolic_multivariate_probability import (ExpectationMatrix, + VarianceMatrix, CrossCovarianceMatrix) + +j, k = symbols("j,k") + +A = MatrixSymbol("A", k, k) +B = MatrixSymbol("B", k, k) +C = MatrixSymbol("C", k, k) +D = MatrixSymbol("D", k, k) + +a = MatrixSymbol("a", k, 1) +b = MatrixSymbol("b", k, 1) + +A2 = MatrixSymbol("A2", 2, 2) +B2 = MatrixSymbol("B2", 2, 2) + +X = RandomMatrixSymbol("X", k, 1) +Y = RandomMatrixSymbol("Y", k, 1) +Z = RandomMatrixSymbol("Z", k, 1) +W = RandomMatrixSymbol("W", k, 1) + +R = RandomMatrixSymbol("R", k, k) + +X2 = RandomMatrixSymbol("X2", 2, 1) + +normal = Normal("normal", 0, 1) + +m1 = Matrix([ + [1, j*Normal("normal2", 2, 1)], + [normal, 0] +]) + +def test_multivariate_expectation(): + expr = Expectation(a) + assert expr == Expectation(a) == ExpectationMatrix(a) + assert expr.expand() == a + + expr = Expectation(X) + assert expr == Expectation(X) == ExpectationMatrix(X) + assert expr.shape == (k, 1) + assert expr.rows == k + assert expr.cols == 1 + assert isinstance(expr, ExpectationMatrix) + + expr = Expectation(A*X + b) + assert expr == ExpectationMatrix(A*X + b) + assert expr.expand() == A*ExpectationMatrix(X) + b + assert isinstance(expr, ExpectationMatrix) + assert expr.shape == (k, 1) + + expr = Expectation(m1*X2) + assert expr.expand() == expr + + expr = Expectation(A2*m1*B2*X2) + assert expr.args[0].args == (A2, m1, B2, X2) + assert expr.expand() == A2*ExpectationMatrix(m1*B2*X2) + + expr = Expectation((X + Y)*(X - Y).T) + assert expr.expand() == ExpectationMatrix(X*X.T) - ExpectationMatrix(X*Y.T) +\ + ExpectationMatrix(Y*X.T) - ExpectationMatrix(Y*Y.T) + + expr = Expectation(A*X + B*Y) + assert expr.expand() == A*ExpectationMatrix(X) + B*ExpectationMatrix(Y) + + assert Expectation(m1).doit() == Matrix([[1, 2*j], [0, 0]]) + + x1 = Matrix([ + [Normal('N11', 11, 1), Normal('N12', 12, 1)], + [Normal('N21', 21, 1), Normal('N22', 22, 1)] + ]) + x2 = Matrix([ + [Normal('M11', 1, 1), Normal('M12', 2, 1)], + [Normal('M21', 3, 1), Normal('M22', 4, 1)] + ]) + + assert Expectation(Expectation(x1 + x2)).doit(deep=False) == ExpectationMatrix(x1 + x2) + assert Expectation(Expectation(x1 + x2)).doit() == Matrix([[12, 14], [24, 26]]) + + +def test_multivariate_variance(): + raises(ShapeError, lambda: Variance(A)) + + expr = Variance(a) + assert expr == Variance(a) == VarianceMatrix(a) + assert expr.expand() == ZeroMatrix(k, k) + expr = Variance(a.T) + assert expr == Variance(a.T) == VarianceMatrix(a.T) + assert expr.expand() == ZeroMatrix(k, k) + + expr = Variance(X) + assert expr == Variance(X) == VarianceMatrix(X) + assert expr.shape == (k, k) + assert expr.rows == k + assert expr.cols == k + assert isinstance(expr, VarianceMatrix) + + expr = Variance(A*X) + assert expr == VarianceMatrix(A*X) + assert expr.expand() == A*VarianceMatrix(X)*A.T + assert isinstance(expr, VarianceMatrix) + assert expr.shape == (k, k) + + expr = Variance(A*B*X) + assert expr.expand() == A*B*VarianceMatrix(X)*B.T*A.T + + expr = Variance(m1*X2) + assert expr.expand() == expr + + expr = Variance(A2*m1*B2*X2) + assert expr.args[0].args == (A2, m1, B2, X2) + assert expr.expand() == expr + + expr = Variance(A*X + B*Y) + assert expr.expand() == 2*A*CrossCovarianceMatrix(X, Y)*B.T +\ + A*VarianceMatrix(X)*A.T + B*VarianceMatrix(Y)*B.T + +def test_multivariate_crosscovariance(): + raises(ShapeError, lambda: Covariance(X, Y.T)) + raises(ShapeError, lambda: Covariance(X, A)) + + + expr = Covariance(a.T, b.T) + assert expr.shape == (1, 1) + assert expr.expand() == ZeroMatrix(1, 1) + + expr = Covariance(a, b) + assert expr == Covariance(a, b) == CrossCovarianceMatrix(a, b) + assert expr.expand() == ZeroMatrix(k, k) + assert expr.shape == (k, k) + assert expr.rows == k + assert expr.cols == k + assert isinstance(expr, CrossCovarianceMatrix) + + expr = Covariance(A*X + a, b) + assert expr.expand() == ZeroMatrix(k, k) + + expr = Covariance(X, Y) + assert isinstance(expr, CrossCovarianceMatrix) + assert expr.expand() == expr + + expr = Covariance(X, X) + assert isinstance(expr, CrossCovarianceMatrix) + assert expr.expand() == VarianceMatrix(X) + + expr = Covariance(X + Y, Z) + assert isinstance(expr, CrossCovarianceMatrix) + assert expr.expand() == CrossCovarianceMatrix(X, Z) + CrossCovarianceMatrix(Y, Z) + + expr = Covariance(A*X, Y) + assert isinstance(expr, CrossCovarianceMatrix) + assert expr.expand() == A*CrossCovarianceMatrix(X, Y) + + expr = Covariance(X, B*Y) + assert isinstance(expr, CrossCovarianceMatrix) + assert expr.expand() == CrossCovarianceMatrix(X, Y)*B.T + + expr = Covariance(A*X + a, B.T*Y + b) + assert isinstance(expr, CrossCovarianceMatrix) + assert expr.expand() == A*CrossCovarianceMatrix(X, Y)*B + + expr = Covariance(A*X + B*Y + a, C.T*Z + D.T*W + b) + assert isinstance(expr, CrossCovarianceMatrix) + assert expr.expand() == A*CrossCovarianceMatrix(X, W)*D + A*CrossCovarianceMatrix(X, Z)*C \ + + B*CrossCovarianceMatrix(Y, W)*D + B*CrossCovarianceMatrix(Y, Z)*C diff --git a/lib/python3.12/site-packages/sympy/stats/tests/test_symbolic_probability.py b/lib/python3.12/site-packages/sympy/stats/tests/test_symbolic_probability.py new file mode 100644 index 0000000000000000000000000000000000000000..edac942ac081c0d44cafd31761b77bc577b6a3fd --- /dev/null +++ b/lib/python3.12/site-packages/sympy/stats/tests/test_symbolic_probability.py @@ -0,0 +1,175 @@ +from sympy.concrete.summations import Sum +from sympy.core.mul import Mul +from sympy.core.numbers import (oo, pi) +from sympy.core.relational import Eq +from sympy.core.symbol import (Dummy, symbols) +from sympy.functions.elementary.exponential import exp +from sympy.functions.elementary.miscellaneous import sqrt +from sympy.functions.elementary.trigonometric import sin +from sympy.integrals.integrals import Integral +from sympy.core.expr import unchanged +from sympy.stats import (Normal, Poisson, variance, Covariance, Variance, + Probability, Expectation, Moment, CentralMoment) +from sympy.stats.rv import probability, expectation + + +def test_literal_probability(): + X = Normal('X', 2, 3) + Y = Normal('Y', 3, 4) + Z = Poisson('Z', 4) + W = Poisson('W', 3) + x = symbols('x', real=True) + y, w, z = symbols('y, w, z') + + assert Probability(X > 0).evaluate_integral() == probability(X > 0) + assert Probability(X > x).evaluate_integral() == probability(X > x) + assert Probability(X > 0).rewrite(Integral).doit() == probability(X > 0) + assert Probability(X > x).rewrite(Integral).doit() == probability(X > x) + + assert Expectation(X).evaluate_integral() == expectation(X) + assert Expectation(X).rewrite(Integral).doit() == expectation(X) + assert Expectation(X**2).evaluate_integral() == expectation(X**2) + assert Expectation(x*X).args == (x*X,) + assert Expectation(x*X).expand() == x*Expectation(X) + assert Expectation(2*X + 3*Y + z*X*Y).expand() == 2*Expectation(X) + 3*Expectation(Y) + z*Expectation(X*Y) + assert Expectation(2*X + 3*Y + z*X*Y).args == (2*X + 3*Y + z*X*Y,) + assert Expectation(sin(X)) == Expectation(sin(X)).expand() + assert Expectation(2*x*sin(X)*Y + y*X**2 + z*X*Y).expand() == 2*x*Expectation(sin(X)*Y) \ + + y*Expectation(X**2) + z*Expectation(X*Y) + assert Expectation(X + Y).expand() == Expectation(X) + Expectation(Y) + assert Expectation((X + Y)*(X - Y)).expand() == Expectation(X**2) - Expectation(Y**2) + assert Expectation((X + Y)*(X - Y)).expand().doit() == -12 + assert Expectation(X + Y, evaluate=True).doit() == 5 + assert Expectation(X + Expectation(Y)).doit() == 5 + assert Expectation(X + Expectation(Y)).doit(deep=False) == 2 + Expectation(Expectation(Y)) + assert Expectation(X + Expectation(Y + Expectation(2*X))).doit(deep=False) == 2 \ + + Expectation(Expectation(Y + Expectation(2*X))) + assert Expectation(X + Expectation(Y + Expectation(2*X))).doit() == 9 + assert Expectation(Expectation(2*X)).doit() == 4 + assert Expectation(Expectation(2*X)).doit(deep=False) == Expectation(2*X) + assert Expectation(4*Expectation(2*X)).doit(deep=False) == 4*Expectation(2*X) + assert Expectation((X + Y)**3).expand() == 3*Expectation(X*Y**2) +\ + 3*Expectation(X**2*Y) + Expectation(X**3) + Expectation(Y**3) + assert Expectation((X - Y)**3).expand() == 3*Expectation(X*Y**2) -\ + 3*Expectation(X**2*Y) + Expectation(X**3) - Expectation(Y**3) + assert Expectation((X - Y)**2).expand() == -2*Expectation(X*Y) +\ + Expectation(X**2) + Expectation(Y**2) + + assert Variance(w).args == (w,) + assert Variance(w).expand() == 0 + assert Variance(X).evaluate_integral() == Variance(X).rewrite(Integral).doit() == variance(X) + assert Variance(X + z).args == (X + z,) + assert Variance(X + z).expand() == Variance(X) + assert Variance(X*Y).args == (Mul(X, Y),) + assert type(Variance(X*Y)) == Variance + assert Variance(z*X).expand() == z**2*Variance(X) + assert Variance(X + Y).expand() == Variance(X) + Variance(Y) + 2*Covariance(X, Y) + assert Variance(X + Y + Z + W).expand() == (Variance(X) + Variance(Y) + Variance(Z) + Variance(W) + + 2 * Covariance(X, Y) + 2 * Covariance(X, Z) + 2 * Covariance(X, W) + + 2 * Covariance(Y, Z) + 2 * Covariance(Y, W) + 2 * Covariance(W, Z)) + assert Variance(X**2).evaluate_integral() == variance(X**2) + assert unchanged(Variance, X**2) + assert Variance(x*X**2).expand() == x**2*Variance(X**2) + assert Variance(sin(X)).args == (sin(X),) + assert Variance(sin(X)).expand() == Variance(sin(X)) + assert Variance(x*sin(X)).expand() == x**2*Variance(sin(X)) + + assert Covariance(w, z).args == (w, z) + assert Covariance(w, z).expand() == 0 + assert Covariance(X, w).expand() == 0 + assert Covariance(w, X).expand() == 0 + assert Covariance(X, Y).args == (X, Y) + assert type(Covariance(X, Y)) == Covariance + assert Covariance(z*X + 3, Y).expand() == z*Covariance(X, Y) + assert Covariance(X, X).args == (X, X) + assert Covariance(X, X).expand() == Variance(X) + assert Covariance(z*X + 3, w*Y + 4).expand() == w*z*Covariance(X,Y) + assert Covariance(X, Y) == Covariance(Y, X) + assert Covariance(X + Y, Z + W).expand() == Covariance(W, X) + Covariance(W, Y) + Covariance(X, Z) + Covariance(Y, Z) + assert Covariance(x*X + y*Y, z*Z + w*W).expand() == (x*w*Covariance(W, X) + w*y*Covariance(W, Y) + + x*z*Covariance(X, Z) + y*z*Covariance(Y, Z)) + assert Covariance(x*X**2 + y*sin(Y), z*Y*Z**2 + w*W).expand() == (w*x*Covariance(W, X**2) + w*y*Covariance(sin(Y), W) + + x*z*Covariance(Y*Z**2, X**2) + y*z*Covariance(Y*Z**2, sin(Y))) + assert Covariance(X, X**2).expand() == Covariance(X, X**2) + assert Covariance(X, sin(X)).expand() == Covariance(sin(X), X) + assert Covariance(X**2, sin(X)*Y).expand() == Covariance(sin(X)*Y, X**2) + assert Covariance(w, X).evaluate_integral() == 0 + + +def test_probability_rewrite(): + X = Normal('X', 2, 3) + Y = Normal('Y', 3, 4) + Z = Poisson('Z', 4) + W = Poisson('W', 3) + x, y, w, z = symbols('x, y, w, z') + + assert Variance(w).rewrite(Expectation) == 0 + assert Variance(X).rewrite(Expectation) == Expectation(X ** 2) - Expectation(X) ** 2 + assert Variance(X, condition=Y).rewrite(Expectation) == Expectation(X ** 2, Y) - Expectation(X, Y) ** 2 + assert Variance(X, Y) != Expectation(X**2) - Expectation(X)**2 + assert Variance(X + z).rewrite(Expectation) == Expectation((X + z) ** 2) - Expectation(X + z) ** 2 + assert Variance(X * Y).rewrite(Expectation) == Expectation(X ** 2 * Y ** 2) - Expectation(X * Y) ** 2 + + assert Covariance(w, X).rewrite(Expectation) == -w*Expectation(X) + Expectation(w*X) + assert Covariance(X, Y).rewrite(Expectation) == Expectation(X*Y) - Expectation(X)*Expectation(Y) + assert Covariance(X, Y, condition=W).rewrite(Expectation) == Expectation(X * Y, W) - Expectation(X, W) * Expectation(Y, W) + + w, x, z = symbols("W, x, z") + px = Probability(Eq(X, x)) + pz = Probability(Eq(Z, z)) + + assert Expectation(X).rewrite(Probability) == Integral(x*px, (x, -oo, oo)) + assert Expectation(Z).rewrite(Probability) == Sum(z*pz, (z, 0, oo)) + assert Variance(X).rewrite(Probability) == Integral(x**2*px, (x, -oo, oo)) - Integral(x*px, (x, -oo, oo))**2 + assert Variance(Z).rewrite(Probability) == Sum(z**2*pz, (z, 0, oo)) - Sum(z*pz, (z, 0, oo))**2 + assert Covariance(w, X).rewrite(Probability) == \ + -w*Integral(x*Probability(Eq(X, x)), (x, -oo, oo)) + Integral(w*x*Probability(Eq(X, x)), (x, -oo, oo)) + + # To test rewrite as sum function + assert Variance(X).rewrite(Sum) == Variance(X).rewrite(Integral) + assert Expectation(X).rewrite(Sum) == Expectation(X).rewrite(Integral) + + assert Covariance(w, X).rewrite(Sum) == 0 + + assert Covariance(w, X).rewrite(Integral) == 0 + + assert Variance(X, condition=Y).rewrite(Probability) == Integral(x**2*Probability(Eq(X, x), Y), (x, -oo, oo)) - \ + Integral(x*Probability(Eq(X, x), Y), (x, -oo, oo))**2 + + +def test_symbolic_Moment(): + mu = symbols('mu', real=True) + sigma = symbols('sigma', positive=True) + x = symbols('x') + X = Normal('X', mu, sigma) + M = Moment(X, 4, 2) + assert M.rewrite(Expectation) == Expectation((X - 2)**4) + assert M.rewrite(Probability) == Integral((x - 2)**4*Probability(Eq(X, x)), + (x, -oo, oo)) + k = Dummy('k') + expri = Integral(sqrt(2)*(k - 2)**4*exp(-(k - \ + mu)**2/(2*sigma**2))/(2*sqrt(pi)*sigma), (k, -oo, oo)) + assert M.rewrite(Integral).dummy_eq(expri) + assert M.doit() == (mu**4 - 8*mu**3 + 6*mu**2*sigma**2 + \ + 24*mu**2 - 24*mu*sigma**2 - 32*mu + 3*sigma**4 + 24*sigma**2 + 16) + M = Moment(2, 5) + assert M.doit() == 2**5 + + +def test_symbolic_CentralMoment(): + mu = symbols('mu', real=True) + sigma = symbols('sigma', positive=True) + x = symbols('x') + X = Normal('X', mu, sigma) + CM = CentralMoment(X, 6) + assert CM.rewrite(Expectation) == Expectation((X - Expectation(X))**6) + assert CM.rewrite(Probability) == Integral((x - Integral(x*Probability(True), + (x, -oo, oo)))**6*Probability(Eq(X, x)), (x, -oo, oo)) + k = Dummy('k') + expri = Integral(sqrt(2)*(k - Integral(sqrt(2)*k*exp(-(k - \ + mu)**2/(2*sigma**2))/(2*sqrt(pi)*sigma), (k, -oo, oo)))**6*exp(-(k - \ + mu)**2/(2*sigma**2))/(2*sqrt(pi)*sigma), (k, -oo, oo)) + assert CM.rewrite(Integral).dummy_eq(expri) + assert CM.doit().simplify() == 15*sigma**6 + CM = Moment(5, 5) + assert CM.doit() == 5**5