Also save/load body-properties. distance_traveled is static method. Bugfix:...

Also save/load body-properties. distance_traveled is static method. Bugfix: when checking if muscles/legs are already set. Bugfix: record final values when record_level==0. Bugfix: pass kwargs to __init__ in random_individual.

manduca/Manducas.py

@@ -19,9 +19,11 @@ class ManducaDeformityError(Exception):
 ManducaBodyProperties = namedtuple('ManducaBodyProperties', ['L0', 'k', 'c', 'm'])
 
 class Manduca(object):
+    """ A simple model of a Manduca Sexta consiting of legs and muscles, modeled as masses and
+    springs, respectively."""
     default_body_properties = ManducaBodyProperties(500, 1, 2, 1)
     damped_body_properties = ManducaBodyProperties(500, 1, 3, 1)
-    def __init__(self, legs, muscles, time_step, fitness_function=None, body_properties=None):
+    def __init__(self, legs, muscles, time_step, body_properties=None, fitness_function=None):
         self._legs = self._muscles = None
         self.legs = np.copy(legs)
         self.muscles = np.copy(muscles)
@@ -29,7 +31,7 @@ class Manduca(object):
         self._fitness = None
         self.check_consistancy()
         if fitness_function == None:
-            self._fitness_function = self.distance_traveled
+            self._fitness_function = Manduca.distance_traveled
         else:
             self._fitness_function = fitness_function
         if body_properties == None:
@@ -53,7 +55,7 @@ class Manduca(object):
     @legs.setter
     def legs(self, legs):
         new_legs = np.array(legs)
-        if self._legs:
+        if self._legs is not None:
             if new_legs.shape != self._legs.shape:
                 msg = 'Ill-defined leg shape. Expected {} but got {}'
                 e = msg.format(self._legs.shape, new_legs.shape)
@@ -68,7 +70,7 @@ class Manduca(object):
     @muscles.setter
     def muscles(self, muscles):
         new_muscles = np.array(muscles)
-        if self._muscles:
+        if self._muscles is not None:
             if new_muscles.shape != self._muscles.shape:
                 msg = 'Ill-defined muscles shape. Expected {} but got {}'
                 e = msg.format(self._muscles.shape, new_muscles.shape)
@@ -102,8 +104,9 @@ class Manduca(object):
             raise(ManducaDeformityError(e))
         return leg_ts
 
-    def distance_traveled(self):
-        positions, velocites, times = self.simulate(record_level=2)
+    @staticmethod
+    def distance_traveled(manduca):
+        positions, velocites, times = manduca.simulate(record_level=0)
         mean_starting_loc = np.mean(positions[0])
         mean_end_loc = np.mean(positions[-1])
         distance_traveled = mean_end_loc - mean_starting_loc
@@ -112,12 +115,7 @@ class Manduca(object):
     @property
     def fitness(self):
         if (self._fitness is None):
-            if isinstance(self._fitness_function, types.MethodType) or \
-               isinstance(self._fitness_function, types.UnboundMethodType):
-                self._fitness = self._fitness_function()
-            elif isinstance(self._fitness_function, types.LambdaType) or \
-                 isinstance(self._fitness_function, types.FunctionType):
-                self._fitness = self._fitness_function(self)
+            self._fitness = self._fitness_function(self)
         return (self._fitness)
 
     @property
@@ -134,7 +132,8 @@ class Manduca(object):
 
     def __repr__(self):
         output = StringIO.StringIO()
-        output.write("\n   legs   |     muscles")
+        output.write("Body Properties: {}\n".format(self.body_properties))
+        output.write("   legs   |     muscles")
         for leg_t, musc_t in zip(self.legs, self.muscles):
             leg_str = " ".join("{:>1}".format(int(l)) for l in leg_t)
             musc_str = " ".join("{:>3}".format(int(m)) for m in musc_t)
@@ -146,12 +145,14 @@ class Manduca(object):
     def __eq__(self, other):
         return (self.time_step == other.time_step and
                 np.array_equal(self.legs, other.legs) and
-                np.array_equal(self.muscles, other.muscles))
+                np.array_equal(self.muscles, other.muscles) and
+                self.body_properties == other.body_properties)
 
     def __ne__(self, other):
         return (self.time_step != other.time_step or
                 not(np.array_equal(self.legs, other.legs)) or
-                not(np.array_equal(self.muscles, other.muscles)))
+                not(np.array_equal(self.muscles, other.muscles) or
+                self.body_properties != other.body_properties))
 
     def __hash__(self):
         leg_str, m_str = self.legs.tostring(), self.muscles.tostring()
@@ -159,12 +160,14 @@ class Manduca(object):
         manduca_str = leg_str + m_str + t_str
         return hash(manduca_str)
 
-    def simulate(self, record_level=0, verbose_record=True, num_sub_intervals = 10):
+    def simulate(self, record_level=0, verbose_record=True, sub_intervals = 10):
         """ record_level: Which intervals to record the position/velocity/time
                           0 -> initial/final only
                           1 -> once for each interval
                           2 -> every sub-interval
                           * -> default to 0
+
+            returns recored  positions, velocities, and times
         """
 
         L0, k, c, m = self.body_properties
@@ -244,7 +247,7 @@ class Manduca(object):
         for time, current_legs, current_muscles in zip(times, self.legs, self.muscles):
             end_time = time + self.time_step
             # Split the interval up into the desired number of sub-intervals
-            interval_times = np.linspace(time, end_time, num_sub_intervals+1)
+            interval_times = np.linspace(time, end_time, sub_intervals+1)
 
             # Concatenate x and v to send to compute slopes
             loc_and_vel = np.hstack ((position, velocities))
@@ -269,6 +272,11 @@ class Manduca(object):
                 recorded_positions.extend(new_positions)
                 recorded_velocities.extend(new_velocities)
                 recorded_times.extend(new_times)
+        # If the record_level is 0, save the final interval as wel
+        if record_level not in [1,2]:
+            recorded_positions.append(position)
+            recorded_velocities.append(velocities)
+            recorded_times.append(end_time)
 
         return recorded_positions, recorded_velocities, recorded_times
 
@@ -278,10 +286,12 @@ class Manduca(object):
         else:
             save_fn = np.savez
         save_fn(file_name, legs=self.legs, muscles=self.muscles,
-                time_step=self.time_step, fitness=self._fitness)
+                time_step=self.time_step, fitness=self._fitness,
+                body_properties=self.body_properties)
 
     @classmethod
     def load(cls, file_name, *other_params, **kwargs):
+        kwargs_copy = dict(kwargs)
         try:
             results = np.load(file_name)
             legs = results['legs']
@@ -290,7 +300,9 @@ class Manduca(object):
             args = []
             for param in other_params:
                 args.append(results[param])
-            new_manduca = cls(legs, muscles, time_step, *args, **kwargs)
+            if 'body_properties' not in kwargs_copy or kwargs_copy['body_properties'] is None:
+                kwargs_copy['body_properties'] = ManducaBodyProperties(*results['body_properties'])
+            new_manduca = cls(legs, muscles, time_step, *args, **kwargs_copy)
             if 'fitness' in results:
                 new_manduca._fitness = np.asscalar(results['fitness'])
             return new_manduca
@@ -300,13 +312,15 @@ class Manduca(object):
             raise(e)
 
     def clone(self):
-        return (self.__class__(np.copy(self.legs), np.copy(self.muscles)))
+        return self.__class__(np.copy(self.legs), np.copy(self.muscles),
+                               time_step=self.time_step,
+                               body_properties=self.body_properties)
 
 
 class SimpleManduca(Manduca):
     """A Manduca that has legs and muscles that are either ON/OFF"""
-    def __init__(self, legs, muscles, time_step, muscle_strength=None):
-        super(SimpleManduca, self).__init__(legs, muscles, time_step)
+    def __init__(self, legs, muscles, time_step, muscle_strength=None, **kwargs):
+        super(SimpleManduca, self).__init__(legs, muscles, time_step, **kwargs)
         locked_legs, unlocked_legs = self.legs == 1, self.legs == 0
         assert (locked_legs | unlocked_legs).all(), "Legs must be 0 or 1"
         if muscle_strength is None:
@@ -325,4 +339,4 @@ class SimpleManduca(Manduca):
         muscles = rng.choice([0, muscle_strength],
                              size=(time_segments, num_legs-1))
         return SimpleManduca(legs, muscles, time_step,
-                             muscle_strength=muscle_strength)
+                             muscle_strength=muscle_strength, **kwargs)