Overload PP ops and use them

earth_model/earth_model.py

@@ -88,41 +88,40 @@ def __init__(self, breakpoints=_bps, density_params=_density_params,
         r2_params = np.tile(np.array([0.0, 0.0, 1000.0**3]),
                             (breakpoints.size - 1, 1))
         self.r2_poly = pp.PeicewisePolynomial(r2_params, breakpoints)
-        self.mass_poly = self.density_poly.mult(self.r2_poly)
+        self.mass_poly = self.density_poly * self.r2_poly
         self.mass_poly = self.mass_poly.antiderivative()
         #  integrating this gives mass:
-        self.mass_poly.coeffs = self.mass_poly.coeffs * 4.0 * np.pi
+        self.mass_poly = self.mass_poly * 4.0 * np.pi
 
         # setup polynomials for MOI. This is 2/3*4*pi*\int rho(r)*r^4 dr
         r4_params = np.tile(np.array([0.0,  0.0, 0.0, 0.0, 1000.0**5]),
                             (breakpoints.size - 1, 1))
         r4_poly = pp.PeicewisePolynomial(r4_params, breakpoints)
-        self.moi_poly = self.density_poly.mult(r4_poly)
+        self.moi_poly = self.density_poly * r4_poly
         self.moi_poly = self.moi_poly.antiderivative()
         #   integrating this gives MOI:
-        self.moi_poly.coeffs = self.moi_poly.coeffs * 4.0 * (2/3) * np.pi
+        self.moi_poly = self.moi_poly * 4.0 * (2/3) * np.pi
 
         # Setup polynomial for gravity
         G = 6.6743E-11
-        gravity_poly = self.density_poly.mult(self.r2_poly)
+        gravity_poly = self.density_poly * self.r2_poly
         # evaluate this to get int(rho.r^2 dr)
         gravity_poly = gravity_poly.integrating_poly()
         # constants outside integral
-        gravity_poly.coeffs = gravity_poly.coeffs * 4.0 * np.pi * G
+        gravity_poly = gravity_poly * 4.0 * np.pi * G
         over_r_sq_poly = pp.PeicewisePolynomial(
             np.zeros((breakpoints.size - 1, 1)), breakpoints,
             np.zeros((breakpoints.size - 1, 3)))
         over_r_sq_poly.negative_coeffs[:, 2] = 1.0/1000.0**2
-        gravity_poly = gravity_poly.mult(over_r_sq_poly)  # Mult by 1/r^2
+        gravity_poly = gravity_poly * over_r_sq_poly  # Mult by 1/r^2
         self.gravity_poly = gravity_poly  # Evaluate to get gravity at r
 
         # Setup polynomial for pressure:
         # integrate from r to r_earth to get pressure
-        self.pressure_poly = self.gravity_poly.mult(self.density_poly)
+        self.pressure_poly = self.gravity_poly * self.density_poly
         self.pressure_poly = self.pressure_poly.antiderivative()
         # Pressure units (/1E9) and density units (*1000.0)
-        self.pressure_poly.coeffs *= 1000.0 / 1.0E9
-        self.pressure_poly.negative_coeffs *= 1000.0 / 1.0E9
+        self.pressure_poly = self.pressure_poly * (1000.0 / 1.0E9)
 
     def density(self, r, break_down=False):
         """

earth_model/peice_poly.py

@@ -217,9 +217,17 @@ def integrating_poly(self):
         return PeicewisePolynomial(ip_coeffs, antiderivative.breakpoints)
 
     def mult(self, other):
-        # FIXME - for this approach brakepoints need to be same place too
+        """
+        Multiplication of peicewise polynomials
+
+        Returns a new PeicewisePolynomial equal to the product of two
+        existing PeicewisePolynomials.
+
+        Note that in the current implementation the breakpoints must
+        match.
+        """
         assert self.coeffs.shape[0] == other.coeffs.shape[0], \
-                                     'different number of breakpoints'
+            'different number of breakpoints'
         mult_breakpoints = self.breakpoints
         mult_coefs = np.zeros((self.coeffs.shape[0],
                                self.coeffs.shape[1]+other.coeffs.shape[1]-1))
@@ -289,3 +297,43 @@ def mult(self, other):
         mult_poly = PeicewisePolynomial(mult_coefs, mult_breakpoints,
                                         mult_negative_coefs)
         return mult_poly
+
+    def scalar_mult(self, scalar):
+        """
+        Multiplication of a peicewise polynomial and a scalar
+
+        Returns a new PeicewisePolynomal resulting from the multiplication
+        of a peicewise polynomial by a scalar (real or int).
+        """
+        result_coeffs = np.copy(self.coeffs) * scalar
+        result_bps = np.copy(self.breakpoints)
+        if self.negative_coeffs is None:
+            result_neg_coeffs = None
+        else:
+            result_neg_coeffs = np.copy(self.negative_coeffs) * scalar
+
+        result_poly = PeicewisePolynomial(result_coeffs, result_bps,
+                                          result_neg_coeffs)
+        return result_poly
+
+    def __mul__(self, other):
+        """
+        Overload * 
+
+        Returns a new PeicewisePolynomal resulting from the 
+        multiplication
+        """
+        if isinstance(other, int) or isinstance(other, float):
+            result = self.scalar_mult(other)
+
+        if isinstance(other, PeicewisePolynomial):
+            result = self.mult(other)
+
+        return result
+
+    def __rmul__(self, other):
+        """
+        Overload * for sclar * PP
+        """
+        return self.__mul__(other)
+

tests/test_peice_poly.py

@@ -376,4 +376,96 @@ def test_mult3():
     # backwards
     calc_poly_mult = poly2.mult(poly1)
     npt.assert_allclose(calc_poly_mult.coeffs, expect_poly_mult.coeffs)
+    npt.assert_allclose(calc_poly_mult.negative_coeffs, expect_poly_mult.negative_coeffs)
+
+    # overloaded
+    calc_poly_mult = poly1 * poly2
+    npt.assert_allclose(calc_poly_mult.coeffs, expect_poly_mult.coeffs)
+    npt.assert_allclose(calc_poly_mult.negative_coeffs, expect_poly_mult.negative_coeffs)
+
+    calc_poly_mult = poly2 * poly1
+    npt.assert_allclose(calc_poly_mult.coeffs, expect_poly_mult.coeffs)
+    npt.assert_allclose(calc_poly_mult.negative_coeffs, expect_poly_mult.negative_coeffs)
+
+
+def test_scalar_mult():
+    # Without negative coefficents, mult by 10
+    poly1 = pp.PeicewisePolynomial(np.array([[4.0, 3.0, 2.0],
+                                             [4.0, 3.0, 2.0]]),
+                                   np.array([0.0, 2.0, 4.0]))
+    expect_poly_mult = pp.PeicewisePolynomial(
+        np.array([[40.0, 30.0, 20.0],
+                  [40.0, 30.0, 20.0]]),
+        np.array([0.0, 2.0, 4.0]))
+    calc_poly_mult = poly1.scalar_mult(10.0)
+    npt.assert_allclose(calc_poly_mult.coeffs, expect_poly_mult.coeffs)
+    assert calc_poly_mult.negative_coeffs is None
+    calc_poly_mult = poly1.scalar_mult(10)
+    npt.assert_allclose(calc_poly_mult.coeffs, expect_poly_mult.coeffs)
+    assert calc_poly_mult.negative_coeffs is None
+
+    # With negative coefficents, mult by 10
+    poly1 = pp.PeicewisePolynomial(np.array([[4.0, 3.0, 2.0],
+                                             [4.0, 3.0, 2.0]]),
+                                   np.array([0.0, 2.0, 4.0]),
+                                   np.array([[4.0, 3.0, 2.0],
+                                             [4.0, 3.0, 2.0]]))
+    expect_poly_mult = pp.PeicewisePolynomial(
+        np.array([[40.0, 30.0, 20.0],
+                  [40.0, 30.0, 20.0]]),
+        np.array([0.0, 2.0, 4.0]),
+        np.array([[40.0, 30.0, 20.0],
+                  [40.0, 30.0, 20.0]]))
+    calc_poly_mult = poly1.scalar_mult(10.0)
+    npt.assert_allclose(calc_poly_mult.coeffs, expect_poly_mult.coeffs)
+    npt.assert_allclose(calc_poly_mult.negative_coeffs, expect_poly_mult.negative_coeffs)
+    calc_poly_mult = poly1.scalar_mult(10)
+    npt.assert_allclose(calc_poly_mult.negative_coeffs, expect_poly_mult.negative_coeffs)
+
+def test_scalar_mult_overload():
+    # Without negative coefficents, mult by 10
+    poly1 = pp.PeicewisePolynomial(np.array([[4.0, 3.0, 2.0],
+                                             [4.0, 3.0, 2.0]]),
+                                   np.array([0.0, 2.0, 4.0]))
+    expect_poly_mult = pp.PeicewisePolynomial(
+        np.array([[40.0, 30.0, 20.0],
+                  [40.0, 30.0, 20.0]]),
+        np.array([0.0, 2.0, 4.0]))
+    calc_poly_mult = poly1 * 10.0
+    npt.assert_allclose(calc_poly_mult.coeffs, expect_poly_mult.coeffs)
+    assert calc_poly_mult.negative_coeffs is None
+    calc_poly_mult = poly1 * 10
+    npt.assert_allclose(calc_poly_mult.coeffs, expect_poly_mult.coeffs)
+    assert calc_poly_mult.negative_coeffs is None
+
+    calc_poly_mult = 10.0 * poly1
+    npt.assert_allclose(calc_poly_mult.coeffs, expect_poly_mult.coeffs)
+    assert calc_poly_mult.negative_coeffs is None
+    calc_poly_mult = 10 * poly1
+    npt.assert_allclose(calc_poly_mult.coeffs, expect_poly_mult.coeffs)
+    assert calc_poly_mult.negative_coeffs is None
+
+
+    # With negative coefficents, mult by 10
+    poly1 = pp.PeicewisePolynomial(np.array([[4.0, 3.0, 2.0],
+                                             [4.0, 3.0, 2.0]]),
+                                   np.array([0.0, 2.0, 4.0]),
+                                   np.array([[4.0, 3.0, 2.0],
+                                             [4.0, 3.0, 2.0]]))
+    expect_poly_mult = pp.PeicewisePolynomial(
+        np.array([[40.0, 30.0, 20.0],
+                  [40.0, 30.0, 20.0]]),
+        np.array([0.0, 2.0, 4.0]),
+        np.array([[40.0, 30.0, 20.0],
+                  [40.0, 30.0, 20.0]]))
+    calc_poly_mult = poly1 * 10.0
+    npt.assert_allclose(calc_poly_mult.coeffs, expect_poly_mult.coeffs)
+    npt.assert_allclose(calc_poly_mult.negative_coeffs, expect_poly_mult.negative_coeffs)
+    calc_poly_mult = poly1 * 10
+    npt.assert_allclose(calc_poly_mult.negative_coeffs, expect_poly_mult.negative_coeffs)
+
+    calc_poly_mult = 10.0 * poly1
+    npt.assert_allclose(calc_poly_mult.coeffs, expect_poly_mult.coeffs)
+    npt.assert_allclose(calc_poly_mult.negative_coeffs, expect_poly_mult.negative_coeffs)
+    calc_poly_mult = 10 * poly1
     npt.assert_allclose(calc_poly_mult.negative_coeffs, expect_poly_mult.negative_coeffs)