pydy · chrisdembia · Aug 7, 2014 · Jul 29, 2014 · Jul 29, 2014 · Jul 30, 2014
diff --git a/examples/double_pendulum/double_pendulum.py b/examples/double_pendulum/double_pendulum.py
@@ -33,11 +33,11 @@
 KM = KanesMethod(N, q_ind=[q1, q2], u_ind=[u1, u2], kd_eqs=kd)
 
 
-(fr, frstar) = KM.kanes_equations(FL, BL)
-kdd = KM.kindiffdict()
-mass_matrix = KM.mass_matrix_full
-forcing_vector = KM.forcing_full
-qudots = mass_matrix.inv() * forcing_vector
-qudots = qudots.subs(kdd)
-qudots.simplify()
-
+KM.kanes_equations(FL, BL)
+#kdd = KM.kindiffdict()
+#mass_matrix = KM.mass_matrix_full
+#forcing_vector = KM.forcing_full
+#qudots = mass_matrix.inv() * forcing_vector
+#qudots = qudots.subs(kdd)
+#qudots.simplify()
+#
diff --git a/examples/double_pendulum/simulate.py b/examples/double_pendulum/simulate.py
@@ -9,44 +9,17 @@
 
 from double_pendulum import *
 
-# List the symbolic arguments
-# ===========================
-
-# Constants
-# ---------
-
 constants = {l: 10.0, m: 10.0, g: 9.81}
 
-# Time-varying
-# ------------
-
-coordinates = [q1, q2]
-
-speeds = [u1, u2]
-
-
-# Generate function that returns state derivatives
-# ================================================
-
-xdot_function = generate_ode_function(mass_matrix, forcing_vector,
-        constants.keys(), coordinates, speeds)
+initial_conditions = {q1: 1.0, q2: 0.0, u1: 0.0, u2: 0.0}
 
-
-# Specify numerical quantities
-# ============================
-
-initial_coordinates = [1.0, 0.0]
-initial_speeds = [0.0, 0.0]
-x0 = concatenate((initial_coordinates, initial_speeds), axis=1)
-
-args = {'constants': constants.values()}
-
-
-# Simulate
-# ========
+sys = System(KM, constants=constants,
+        initial_conditions=initial_conditions)
 
 frames_per_sec = 60
 final_time = 5.0
 
 t = linspace(0.0, final_time, final_time * frames_per_sec)
-x = odeint(xdot_function, x0, t, args=(args,))
+
+x = sys.integrate(times)
+
diff --git a/examples/double_pendulum/visualize.py b/examples/double_pendulum/visualize.py
@@ -58,7 +58,6 @@
 # Create the visualization
 # ========================
 
-scene.generate_visualization_json(coordinates + speeds, constants.keys(), x,
-        constants.values())
+scene.generate_visualization_json(sys)
 
 scene.display()
diff --git a/pydy/codegen/tests/models.py b/pydy/codegen/tests/models.py
@@ -9,7 +9,7 @@
 import sympy.physics.mechanics as me
 
 
-def generate_mass_spring_damper_equations_of_motion(external_force=True):
+def generate_mass_spring_damper_equations_of_motion(external_force=True, kane=False):
     """Returns the symbolic equations of motion and associated variables for a
     simple one degree of freedom mass, spring, damper system with gravity
     and an optional external specified force.
@@ -90,8 +90,11 @@ def generate_mass_spring_damper_equations_of_motion(external_force=True):
     else:
         specified = None
 
-    return (mass_matrix, forcing_vector, constants, coordinates, speeds,
-            specified)
+    if kane:
+        return kane
+    else:
+        return (mass_matrix, forcing_vector, constants, coordinates, speeds,
+                specified)
 
 
 def generate_n_link_pendulum_on_cart_equations_of_motion(n, cart_force=True,

diff --git a/pydy/system.py b/pydy/system.py
@@ -0,0 +1,105 @@
+
+from .codegen.code import generate_ode_function
+from scipy.integrate import odeint
+
+class System(object):
+    """Manages the simulation (integration) of a system whose equations are
+    given by KanesMethod or LagrangesMethod.
+
+    All attributes can be set directly. With the exception of `method`, the
+    attributes can also be set via keyword arguments to the constructor.
+
+    Attributes
+    ----------
+    method : sympy.physics.mechanics.KanesMethod or
+                sympy.physics.mechanics.LagrangesMethod
+        The method used to generate the equations of motion.
+    constants : dict, optional (default: all 1.0)
+        Numerical values for the constants in the problem. Keys are the symbols
+        for the constants, and values are floats. Constants that are not
+        specified in this dict are given a default value of 1.0.
+    specified_symbols : iterable of symbols, optional
+        The symbols for the quantities whose numerical values you want to
+        specify. The order of these must match the order of the
+        `specified_values`. If not provided, we get all the specified symbols
+        from the `method`, and set their numerical value to 0. You can also
+        provide a subset of symbols for this attribute, in which case the
+        remaining symbols still have their numerical value set to 0.
+    specified_values : iterable of floats or functions, optional
+        The numerical values of the `specified_symbols`, or functions (which
+        take the same arguments as f, above) that can generate the numerical
+        values. If not provided, we get all the specified symbols
+        from the `method`, and set their numerical value to 0. You can also
+        provide a subset of symbols for this attribute, in which case the
+        remaining symbols still have their numerical value set to 0.
+    code_gen_backend : str, optional (default: 'lambdify')
+        The backend used to generate the ode function.
+        See the documentation of pydy.codegen.code.generate_ode_function`.
+    ode_solver : function, optional (default: scipy.integrate.odeint)
+        A function that performs forward integration. It must have the same
+        signature as odeint, which is::
+
+            x_history = ode_solver(f, x0, t, args=(args,))
+
+        where f is a function f(x, t), x0 are the initial conditions, x_history
+        is the history, and args is a keyword argument takes arguments that are
+        then passed to f. TODO
+    initial_conditions : dict, optional (default: all zero)
+        Initial conditions for all coordinates and speeds. Keys are the symbols
+        for the coordinates and speeds, and values are floats. Coordinates or
+        speeds that are not specified in this dict are given a default value of
+        zero.
+
+    """
+    def __init__(self, method, **kwargs):
+        self.method = method
+        for k, v in kwargs:
+            setattr(self, k, v)
+        if self.constants != None:
+            self.constants = self._find_constants()
+        if self.specified_symbols != None and self.specified_values != None:
+            self._check_specified()
+            # TODO
+            pass
+        if self.code_gen_backend != None:
+            self.code_gen_backend = 'lambdify'
+        if ode_solver != None:
+            self.ode_solver = odeint
+
+    def _find_constants(self):
+        othersymbols = method._find_othersymbols()
+        return zip(othersymbols, len(othersymbols) * [1.0])
+
+    def _find_specified_symbols(self):
+        specified_symbols = method._find_dynamicsymbols()
+        return zip
+
+    def generate_ode_function(self):
+        """Calls `pydy. TODO
+
+        """
+        self.rhs = generate_ode_function(
+                self.method.mass_matrix_full, self.method.forcing_full,
+                self.constants.keys(),
+                self.method._q, self.method._u)
+
+
+
+#for k, v in kwargs:
+#    setattr(self, k, v)
+#sys.generate_ode_function(backend=...)
+#sys.code_gen_backend = 
+#sys.ode_solver = odeint
+#sys.specified_symbols =
+#sys.specified_values =
+#sys.initial_conditions =
+#sys.constants = 
+#
+#class Sys:
+#    self.initial_conditions = zero(
+#    def integrate(self, times):
+#        # make rhs
+#        sys.ode_solver(self.rhs, self.initial_conditions, times,
+#                args=(self.con)
+#    
+
diff --git a/pydy/tests/__init__.py b/pydy/tests/__init__.py
diff --git a/pydy/tests/test_system.py b/pydy/tests/test_system.py
@@ -0,0 +1,116 @@
+
+import numpy as np
+from numpy import testing
+from sympy import symbols
+from sympy.physics.mechanics import dynamicsymbols
+from scipy.integrate import odeint
+
+from ..system import System
+from ..codegen.tests.models import \
+    generate_mass_spring_damper_equations_of_motion as mass_spring_damper
+
+
+class TestSystem():
+
+    def setup(self):
+
+        self.kane = mass_spring_damper(kane=True)
+        self.specified_symbol = dynamicsymbols('F')
+        self.constant_map = dict(zip(symbols('m, k, c, g'),
+                                     [2.0, 1.5, 0.5, 9.8]))
+        self.sys = System(self.kane,
+                          code_gen_backend='lambdify',
+                          ode_solver=odeint,
+                          specified_symbols=(self.specified_symbol,),
+                          specified_values=np.ones(1),
+                          initial_conditions=np.zeros(2),
+                          constants=self.constant_map)
+
+    def test_init(self):
+
+        sys = System(self.kane)
+
+        assert sys.method is self.kane
+        assert sys.code_gen_backend == 'lambdify'
+        assert sys.specified_symbols is dynamicsymbols('F')
+        testing.assert_allclose(sys.specified_values, np.zeros(1))
+        testing.assert_allclose(sys.initial_conditions, np.zeros(2))
+        assert sys.constants == dict(zip(symbols('m, k, c, g'),
+                                     [1.0, 1.0, 1.0, 1.0]))
+
+        sys = System(self.kane,
+                     code_gen_backend='lambdify',
+                     ode_solver=odeint,
+                     specified_symbols=(self.specified_symbol,),
+                     specified_values=np.ones(1),
+                     initial_conditions=np.zeros(2),
+                     constants=self.constant_map)
+
+        assert sys.method is self.kane
+        assert sys.code_gen_backend == 'lambdify'
+        assert sys.specified_symbols is dynamicsymbols('F')
+        testing.assert_allclose(sys.specified_values, np.ones(1))
+        testing.assert_allclose(sys.initial_conditions, np.zeros(2))
+        assert sys.constants == self.constant_map
+
+        sys.backend == 'cython'
+
+        # You must set both specified symbols and values.
+        assert_raises(ValueError,
+                      System(self.kane,
+                             specified_symbols=(self.specified_symbol)))
+
+        assert_raises(ValueError,
+                      System(self.kane,
+                             specified_values=(self.specified_symbol)))
+
+        # Check to make sure the specified function returns the correct
+        # number of values.
+        assert_raises(ValueError,
+                      System(self.kane,
+                             specified_symbols=(self.specified_symbol),
+                             specified_values=lambda x, t: np.ones(10))
+
+
+        # The specified symbol must exist in the equations of motion and not
+        # be a state.
+        assert_raises(ValueError,
+                      System(self.kane,
+                             specified_symbols=(dynamicsymbols('G')),
+                             specified_values=np.ones(1)))
+
+        # The same number of specified symbols and values are needed.
+        assert_raises(ValueError,
+                      System(self.kane,
+                             specified_symbols=(self.specified_symbol),
+                             specified_values=np.ones(2))
+
+
+    def test_generate_ode_function(self):
+
+        rhs = self.sys.generate_ode_function()
+
+        assert rhs is self.rhs
+
+        args = {'constants': self.constant_map.values,
+                'specified': np.zeros(1)}
+
+        actual = rhs(np.ones(2), 0.0, args=(args,))
+
+        testing.assert_allclose(actual,
+                                np.array([ , ]))
+
+    def find_constants(self):
+
+        constant_dict = self.sys._find_constants()
+
+        assert constant_dict = self.constant_map
+
+    def find_specified_symbols(self):
+
+        specified_dict = self.sys._find_specified_symbols()
+
+        assert specified_dict = {self.dynamic_symbol: 0.0}
+
+    def test_integrate(self):
+