diff --git a/earth_model/earth_model.py b/earth_model/earth_model.py index 1f7f5ed..2807f70 100644 --- a/earth_model/earth_model.py +++ b/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): """ diff --git a/earth_model/peice_poly.py b/earth_model/peice_poly.py index 93a780f..ab4c7ce 100644 --- a/earth_model/peice_poly.py +++ b/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) + diff --git a/tests/test_peice_poly.py b/tests/test_peice_poly.py index d83d27e..4109857 100644 --- a/tests/test_peice_poly.py +++ b/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) \ No newline at end of file