Merge pull request #59 from JQIamo/v102_bug_fixes

V102 bug fixes

pylcp/fields.py

@@ -42,8 +42,32 @@ def return_constant_val_t(t, val):
     else:
         return val
 
-def promote_to_lambda(val, var_name=None, type='Rt'):
-    if type is 'Rt':
+def promote_to_lambda(val, var_name='', type='Rt'):
+    """
+    Promotes a constant or callable to a lambda function with proper arguments.
+
+    Parameters
+    ----------
+        val : array_like or callable
+            The value to promote.  Can either be a function, array_like
+            (vector), or a scalar (constant).
+        var_name : str, optional
+            Name of the variable attempting to be promoted.  Useful for error
+            messages. Default: empty string.
+        type : str, optional
+            The arguments of the lambda function we are creating.  If `Rt`,
+            the lambda function returned has ``(R, t)`` as its arguments.  If
+            `t`, it has only `t` as its arguments.  Default: `Rt`.
+
+
+    Returns
+    -------
+        func : callable
+            lambda function created by this function
+        sig : string
+            Either `(R,t)` or `(t)`
+    """
+    if type == 'Rt':
         if not callable(val):
             if isinstance(val, list) or isinstance(val, np.ndarray):
                 func = lambda R=np.array([0., 0., 0.]), t=0.: return_constant_vector(R, t, val)
@@ -58,12 +82,15 @@ def promote_to_lambda(val, var_name=None, type='Rt'):
             elif ('(R, t)' in sig or '(r, t)' in sig or '(x, t)' in sig):
                 func = lambda R=np.array([0., 0., 0.]), t=0.: val(R, t)
                 sig = '(R, t)'
+            elif '(t)' in sig:
+                func = lambda R=np.array([0., 0., 0.]), t=0.: val(t)
+                sig = '(R, t)'
             else:
                 raise TypeError('Signature [%s] of function %s not'+
                                 'understood.'% (sig, var_name))
 
         return func, sig
-    elif type is 't':
+    elif type == 't':
         if not callable(val):
             func = lambda t=0.: return_constant_val_t(t, val)
             sig = '()'

pylcp/gratings.py

@@ -63,7 +63,7 @@ def __init__(self, delta=-1., s=1., nr=3, thd=np.pi/4,
                  pol=np.array([-1/np.sqrt(2), 1j/np.sqrt(2), 0]),
                  reflected_pol=np.array([np.pi, 0]),
                  reflected_pol_basis='poincare',
-                 eta=None, grating_angle=0):
+                 eta=None, grating_angle=0, **kwargs):
         """
         Creates beams that would be made from a grating.
         Parameters:
@@ -104,7 +104,7 @@ def __init__(self, delta=-1., s=1., nr=3, thd=np.pi/4,
 
         self.add_laser(infinitePlaneWaveBeam(kvec=np.array([0., 0., 1.]),
                                              pol=pol, s=s, delta=delta,
-                                             pol_coord='cartesian'))
+                                             pol_coord='cartesian', **kwargs))
 
         # Store the input polarization as a Carterian coordiante:
         self.input_pol = self.beam_vector[0].cartesian_pol()
@@ -119,7 +119,8 @@ def __init__(self, delta=-1., s=1., nr=3, thd=np.pi/4,
                                                  pol=pol_refs[:, ii],
                                                  s=self.eta*s/np.cos(self.thd),
                                                  delta=delta,
-                                                 pol_coord='cartesian'))
+                                                 pol_coord='cartesian',
+                                                 **kwargs))
 
     def _calculate_reflected_kvecs_and_pol(self, reflected_pol,
                                            reflected_pol_basis):
@@ -255,7 +256,8 @@ def polarization_ellipse(self, R=np.array([0., 0., 0.]), t=0):
 class inputGaussianBeam(clippedGaussianBeam):
     def __init__(self, kvec=np.array([0, 0, 1]), s=1., delta=-1.,
                  pol=np.array([-1/np.sqrt(2), 1j/np.sqrt(2), 0]),
-                 pol_coord='cartesian', wb=1., rs=2., nr=3, center_hole=0.0, zgrating=1.0, grating_angle=0, **kwargs):
+                 pol_coord='cartesian', wb=1., rs=2., nr=3, center_hole=0.0,
+                 zgrating=1.0, grating_angle=0, **kwargs):
 
         if not np.allclose(kvec[:2], 0.):
             raise ValueError('inputGaussianBeam must be aligned with z axis')
@@ -434,7 +436,10 @@ def __init__(self, delta=-1., s=1., nr=3, thd=np.pi/4,
                  center_hole=0.0,
                  outer_radius=10.0,
                  zgrating=1.0,
-                 grating_angle=0):
+                 grating_angle=0,
+                 input_phase=0.0,
+                 diff_phases=None,
+                 **kwargs):
         """
         Creates beams that would be made from a grating.
         Parameters:
@@ -467,7 +472,10 @@ def __init__(self, delta=-1., s=1., nr=3, thd=np.pi/4,
             center_hole: inscribed radius of center hole.
             outer_radius: outer radius of the diffraction gratings.
             zgrating: z position of the diffraction grating chip.
-            grating_angle: overall azimuthal rotation of the grating
+            grating_angle: overall azimuthal rotation of the grating.
+            input_phase: Phase of input beam.
+            diff_phase: List or numpy array containing phases of diffracted
+            beams.
         """
         # I would like use to do super().super().__init__(), but that does not work.
         # Nonetheless, these are the only two things that need to be done
@@ -484,12 +492,17 @@ def __init__(self, delta=-1., s=1., nr=3, thd=np.pi/4,
         else:
             self.eta = eta
 
+        if diff_phases is None:
+            diff_phases = np.zeros(nr)
+
         self.add_laser(inputGaussianBeam(kvec=np.array([0., 0., 1.]),
                                          pol=pol, s=s, delta=delta,
                                          pol_coord='cartesian', wb=wb, rs=rs,
                                          nr=self.nr, center_hole=center_hole,
                                          zgrating=zgrating,
-                                         grating_angle=self.grating_angle))
+                                         grating_angle=self.grating_angle,
+                                         phase=input_phase,
+                                         **kwargs))
 
         # Store the input polarization as a Carterian coordiante:
         self.input_pol = self.beam_vector[0].cartesian_pol()
@@ -511,7 +524,9 @@ def __init__(self, delta=-1., s=1., nr=3, thd=np.pi/4,
                                                  center_hole=center_hole,
                                                  outer_radius=outer_radius,
                                                  zgrating=zgrating,
-                                                 grating_angle=self.grating_angle))
+                                                 grating_angle=self.grating_angle,
+                                                 phase=diff_phases[ii],
+                                                 **kwargs))
 
         if eta0 is not None:
             raise NotImplementedError("Zeroth-order reflected beam has not been implemented yet.")

pylcp/hamiltonian.py

@@ -168,12 +168,12 @@ def __search_elem_label(self, label):
         return ind
 
     def __make_elem_label(self, type, state_label):
-        if type is 'H_0' or type is 'mu_q':
+        if type == 'H_0' or type == 'mu_q':
             if not isinstance(state_label, str):
                 raise TypeError('For type %s, state label %s must be a' +
                                 ' string.' % (type,state_label))
             return '<%s|%s|%s>' % (state_label, type, state_label)
-        elif type is 'd_q':
+        elif type == 'd_q':
             if not isinstance(state_label, list):
                 raise TypeError('For type %s, state label %s must be a' +
                                 ' list of two strings.' % (type, state_label))

pylcp/integration_tools.py

@@ -1,6 +1,7 @@
 from __future__ import division, print_function, absolute_import
 import inspect
 import numpy as np
+from inspect import signature
 from scipy.integrate._ivp.bdf import BDF
 from scipy.integrate._ivp.radau import Radau
 from scipy.integrate._ivp.rk import RK23, RK45
@@ -34,11 +35,16 @@ class RandomOdeResult(OptimizeResult):
 class parallelIntegrator(object):
     """
     parallelIntegrator: a class to integrate a function as it is being called
-    
+
     Parameters:
     ----------
     func : callable
-        The function that is to be integrated
+        The function that is to be integrated.  It can have the form func(t) or
+        func(t,y).
+
+    y0 : float or array, optional
+        The initial value of y.  Default value is 0.
+
     method : string, optional
         Integration method to use:
             * 'RK45' (default): Explicit Runge-Kutta method of order 5(4) [1]_.
@@ -72,18 +78,24 @@ class parallelIntegrator(object):
               from ODEPACK.
     tmax : float, optional
         Maximum magnitude of the time.  By default, 1e9.
-    kwargs : 
+    kwargs :
         Options passed to a chosen OdeSolver.
-        
+
     Attributes
     ----------
     t0 : initial time of the integrator
     tlast : last time evaluated
     direction : direction of the integrator
     tmax : maximium value of integrator
     """
-    def __init__(self, func, method='RK45', tmax=1e9, **kwargs):
-        self.func = func
+    def __init__(self, func, y0=[0.], method='RK45', tmax=1e9, **kwargs):
+        if '(t, y' in str(signature(func)):
+            self.func = func
+        elif '(t' in str(signature(func)):
+            self.func = lambda t, y: func(t)
+        else:
+            raise ValueError('signature %s for func not recognized'%str(signature(func)))
+
         self.t0 = None
         self.tlast = None
         self.direction = +1
@@ -102,65 +114,83 @@ def __init__(self, func, method='RK45', tmax=1e9, **kwargs):
         elif method == 'LSODA':
             self.intobj = LSODA
         else:
-            raise ValueError('Method %s not recognized'%self.method) 
-
+            raise ValueError('Method %s not recognized'%self.method)
+
+        self.y0 = np.array(y0)
         self.extra_kwargs = kwargs
 
     def __call__(self, t):
         """
         __call: return value at time t:
-        
+
         Parameters:
         -----------
-        t : float
+        t : float or array_like
             time at which to evaluate function
-        
+
         Returns:
         --------
-        y : float
+        y : float or array_like
             value of the function at time t.
         """
+        if isinstance(t, np.ndarray):
+            self.__step(np.amin(t)) # start the intergrator
+            self.__step(np.amax(t)) # step the integrator through to max value
+
+            # Rebuild the solution:
+            sol = OdeSolution(self.ts, self.interpolants)
+
+            return sol(t) # return the full array
+        else:
+            self.__step(t) # step integrator
+
+            if t==self.t0:
+                return self.y0
+            else:
+                # Rebuild the solution:
+                sol = OdeSolution(self.ts, self.interpolants)
+
+                return sol(t)
+
+    def __step(self, t):
         # Is this the first call, or did we return to the initial time?
         if self.t0 is None or t==self.t0:
             self.t0 = t
             self.tlast = None
             self.interpolants = []
             self.ts = [t]
-
-            return 0.
+
+            # Return the initial value:
+            return self.y0
+
         # Second call, we will now establish a direction and create the solver:
-        elif self.tlast == None: 
+        elif self.tlast == None:
             if t>self.t0:
                 self.direction = +1
             elif t<self.t0:
                 self.direction = -1
-            
-            self.integrator=self.intobj(lambda t, y: self.func(t), self.t0, [0.],
+
+            self.integrator=self.intobj(self.func, self.t0, self.y0,
                                         self.direction*self.tmax, **self.extra_kwargs)
-        # Did we go to a value smaller than our initial direction?
+
+        # Did we go to a value smaller than our initial value, given the
+        # direction?
         elif (t<self.t0 and self.direction==+1) or (t>self.t0 and self.direction==-1):
             # Reset the integrator.
             self.t0=None
             self.tlast=None
+            # Cute way to reset the integrator to a new starting t:
             return self(t)
-        
+
         # If we made it here, we did not reset yet:
         self.tlast = t
-
-        if t == self.t0:
-            return 0.
-        else:
-            while self.integrator.t<t:
-                self.integrator.step()
-                sol = self.integrator.dense_output()
-                self.interpolants.append(sol)
-                self.ts.append(self.integrator.t)
-
-            # Rebuild the solution:
-            sol = OdeSolution(self.ts, self.interpolants)
-
-            return sol(t)
 
+        # Now integrator up to t:
+        while self.integrator.t<t:
+            self.integrator.step()
+            sol = self.integrator.dense_output()
+            self.interpolants.append(sol)
+            self.ts.append(self.integrator.t)
 
 def solve_ivp_random(fun, random_func, t_span, y0,  method='RK45', t_eval=None,
                      dense_output=False, events=None, vectorized=False,
@@ -519,7 +549,7 @@ def solve_ivp_random(fun, random_func, t_span, y0,  method='RK45', t_eval=None,
 
     max_step_initial = options.pop('initial_max_step', np.inf)
     max_step_global = options.pop('max_step', np.inf)
-    
+
     solver = method(fun, t0, y0, tf, vectorized=vectorized,
                     max_step=max_step_initial, **options)
 

pylcp/obe.py

@@ -1072,10 +1072,13 @@ def find_equilibrium_force(self, deltat=500, itermax=100, Npts=5001,
         ii : int
             Number of iterations needed to converge.
         """
-        if initial_rho is 'rateeq':
+        if initial_rho == 'rateeq':
             self.set_initial_rho_from_rateeq()
-        elif initial_rho is 'equally':
+        elif initial_rho == 'equally':
             self.set_initial_rho_equally()
+        elif initial_rho == 'frompops':
+            Npop = kwargs.pop('init_pop', None)
+            self.set_initial_rho_from_populations(Npop)
         else:
             raise ValueError('Argument initial_rho=%s not understood'%initial_rho)
 

pylcp/rateeq.py

@@ -485,6 +485,12 @@ def set_initial_pop(self, N0):
             :math:`n` elements, where :math:`n` is the total number of states
             in the system.
         """
+        if len(N0) != self.hamiltonian.n:
+            raise ValueError('Npop has only %d entries for %d states.' %
+                             (len(N0), self.hamiltonian.n))
+        if np.any(np.isnan(N0)) or np.any(np.isinf(N0)):
+            raise ValueError('Npop has NaNs or Infs!')
+
         self.N0 = N0
 
     def set_initial_pop_from_equilibrium(self):
@@ -501,7 +507,7 @@ def evolve_populations(self, t_span, **kwargs):
 
         This function integrates the rate equations to determine how the
         populations evolve in time.  Any initial velocity is kept constant.
-        It is analogous to obe.evolve_densityv().
+        It is analogous to obe.evolve_density().
 
         Parameters
         ----------

setup.py

@@ -11,7 +11,8 @@
     install_requires=["docutils>=0.3",
                       "numpy>=1.18",
                       "numba>=0.48",
-                      "scipy>=1.4.1"],
+                      "scipy>=1.4.1",
+                      "sympy>=1.6"],
     python_requires=">=3.6, <4",
     package_data={
         # If any package contains *.txt or *.rst files, include them: