diff --git a/.github/environment_ci.yml b/.github/environment_ci.yml index cc85be11..0ba145e5 100644 --- a/.github/environment_ci.yml +++ b/.github/environment_ci.yml @@ -22,6 +22,7 @@ dependencies: # testing - pytest - pytest-cov + - sympy # Building - pip diff --git a/.github/workflows/test_with_conda.yml b/.github/workflows/test_with_conda.yml index 784f0339..79e814a1 100644 --- a/.github/workflows/test_with_conda.yml +++ b/.github/workflows/test_with_conda.yml @@ -68,7 +68,7 @@ jobs: - name: "Upload coverage to Codecov" if: matrix.python-version == '3.11' - uses: codecov/codecov-action@v3 + uses: codecov/codecov-action@v4 with: token: ${{ secrets.CODECOV_TOKEN }} fail_ci_if_error: true diff --git a/geoana/earthquake/oksar.py b/geoana/earthquake/oksar.py index 2815e0b6..e0978ab8 100644 --- a/geoana/earthquake/oksar.py +++ b/geoana/earthquake/oksar.py @@ -1657,7 +1657,7 @@ def example(): # data=data # ) - data = np.fromstring(dinar_file.content, np.float32) + data = np.frombuffer(dinar_file.content, np.float32) title = 'Dinar, Turkey' location = [706216.0606, 4269238.9999] location_UTM_zone = 35 diff --git a/geoana/em/base.py b/geoana/em/base.py index 2a611bd0..4c089fc8 100644 --- a/geoana/em/base.py +++ b/geoana/em/base.py @@ -58,6 +58,26 @@ def sigma(self, value): self._sigma = value + @property + def rho(self): + """Electrical resistivity in Ohm m + + Returns + ------- + float + Electrical resistivity in Ohm m + """ + return 1/self.sigma + + @rho.setter + def rho(self, value): + try: + value = float(value) + except: + raise TypeError(f"resistivity must be a number, got {type(value)}") + + self.sigma = 1/value + @property def mu(self): """Magnetic permeability in H/m @@ -177,6 +197,8 @@ def orientation(self, var): var = np.r_[0., 1., 0.] elif var.upper() == 'Z': var = np.r_[0., 0., 1.] + else: + raise ValueError("Orientation must be one of {'X','Y','Z'}") else: try: var = np.asarray(var, dtype=float) @@ -187,7 +209,7 @@ def orientation(self, var): var = np.squeeze(var) if var.shape != (3, ): raise ValueError( - f"orientation must be array_like with shape (3,), got {len(var)}" + f"orientation must be array_like with shape (3,), got {var.shape}" ) # Normalize the orientation diff --git a/geoana/em/fdem/base.py b/geoana/em/fdem/base.py index 467060d0..d8e30b77 100644 --- a/geoana/em/fdem/base.py +++ b/geoana/em/fdem/base.py @@ -203,13 +203,11 @@ def frequency(self, value): value = np.asarray(value, dtype=float) except: raise TypeError(f"frequencies are not a valid type") - value = np.atleast_1d(value) - # Enforce positivity and dimensions if (value < 0.).any(): raise ValueError("All frequencies must be greater than 0") if value.ndim > 1: - raise TypeError(f"frequencies must be ('*') array") + raise TypeError(f"frequencies must have at most 1 dimension.") self._frequency = value diff --git a/geoana/em/fdem/halfspace.py b/geoana/em/fdem/halfspace.py index 6e0a8ebc..6646c17c 100644 --- a/geoana/em/fdem/halfspace.py +++ b/geoana/em/fdem/halfspace.py @@ -34,7 +34,7 @@ def magnetic_field(self, xy, field="secondary"): Returns ------- - (n_freq, ..., 3) numpy.array of complex + (n_freq, ..., 3) numpy.ndarray of complex Magnetic field at all frequencies for the gridded locations provided. Output array is squeezed when n_freq and/or n_loc = 1. diff --git a/geoana/em/fdem/layered.py b/geoana/em/fdem/layered.py index e3f11ed7..bf363adc 100644 --- a/geoana/em/fdem/layered.py +++ b/geoana/em/fdem/layered.py @@ -300,7 +300,7 @@ def magnetic_field(self, xyz, field="secondary"): Returns ------- - (n_freq, n_loc, 3) numpy.array of complex + (n_freq, n_loc, 3) numpy.ndarray of complex Magnetic field at all frequencies for the gridded locations provided. Output array is squeezed when n_freq and/or n_loc = 1. diff --git a/geoana/em/fdem/wholespace.py b/geoana/em/fdem/wholespace.py index 3712c591..d9d0cedb 100644 --- a/geoana/em/fdem/wholespace.py +++ b/geoana/em/fdem/wholespace.py @@ -1,7 +1,6 @@ import numpy as np from geoana.em.fdem.base import BaseFDEM -from geoana.spatial import repeat_scalar -from geoana.utils import check_xyz_dim +from geoana.utils import check_xyz_dim, append_ndim from geoana.em.base import BaseElectricDipole, BaseMagneticDipole @@ -39,7 +38,7 @@ def vector_potential(self, xyz): Returns ------- - (n_freq, ..., 3) numpy.array of complex + (n_freq, ..., 3) numpy.ndarray of complex Magnetic vector potential at all frequencies for the gridded locations provided. Output array is squeezed when n_freq and/or n_loc = 1. @@ -91,13 +90,15 @@ def vector_potential(self, xyz): >>> ax2.autoscale(tight=True) """ - r = np.linalg.norm(xyz - self.location, axis=-1) - k = self.wavenumber - for dim in range(r.ndim): - k = k[:, None] + xyz = check_xyz_dim(xyz) + + r_vec = xyz - self.location + r = np.linalg.norm(r_vec, axis=-1, keepdims=True) + + k = append_ndim(self.wavenumber, r.ndim) a = self.current * self.length / (4*np.pi*r) * np.exp(-1j*k * r) - return (a[..., None] * self.orientation).squeeze() + return a * self.orientation def electric_field(self, xyz): r"""Electric field for the harmonic current dipole at a set of gridded locations. @@ -126,7 +127,7 @@ def electric_field(self, xyz): Returns ------- - (n_freq, ..., 3) numpy.array of complex + (n_freq, ..., 3) numpy.ndarray of complex Electric field at all frequencies for the gridded locations provided. Output array is squeezed when n_freq and/or n_loc = 1. @@ -181,11 +182,12 @@ def electric_field(self, xyz): >>> ax2.set_title('Imag component {} Hz'.format(frequency[f_ind])) """ - dxyz = xyz - self.location - r = np.linalg.norm(dxyz, axis=-1) - k = self.wavenumber - for dim in range(r.ndim): - k = k[:, None] + xyz = check_xyz_dim(xyz) + + r_vec = xyz - self.location + r = np.linalg.norm(r_vec, axis=-1, keepdims=True) + r_hat = r_vec / r + k = append_ndim(self.wavenumber, r.ndim) kr = k * r ikr = 1j * kr @@ -194,12 +196,10 @@ def electric_field(self, xyz): (self.current * self.length) / (4 * np.pi * self.sigma * r**3) * np.exp(-ikr) ) - symmetric_term = ( - ((dxyz @ self.orientation) * (-kr**2 + 3*ikr + 3) / r**2)[..., None] * dxyz - ) - oriented_term = (kr**2 - ikr - 1)[..., None] * self.orientation + symmetric_term = (r_hat @ self.orientation)[..., None] * (-kr**2 + 3*ikr + 3) * r_hat + oriented_term = (kr**2 - ikr - 1) * self.orientation - return (front_term[..., None] * (symmetric_term + oriented_term)).squeeze() + return front_term * (symmetric_term + oriented_term) def current_density(self, xyz): r"""Current density for the harmonic current dipole at a set of gridded locations. @@ -223,12 +223,12 @@ def current_density(self, xyz): Parameters ---------- - xyz : (n, 3) numpy.ndarray + xyz : (..., 3) numpy.ndarray Gridded xyz locations Returns ------- - (n_freq, n_loc, 3) numpy.array of complex + (n_freq, ..., 3) numpy.ndarray of complex Current density at all frequencies for the gridded locations provided. Output array is squeezed when n_freq and/or n_loc = 1. @@ -311,10 +311,9 @@ def magnetic_field(self, xyz): Returns ------- - (n_freq, ..., 3) numpy.array of complex + (n_freq, ..., 3) numpy.ndarray of complex Magnetic field at all frequencies for the gridded - locations provided. Output array is squeezed when n_freq and/or - n_loc = 1. + locations provided. Examples -------- @@ -366,11 +365,13 @@ def magnetic_field(self, xyz): >>> ax2.set_title('Imag component {} Hz'.format(frequency[f_ind])) """ - dxyz = xyz - self.location - r = np.linalg.norm(dxyz, axis=-1) - k = self.wavenumber - for dim in range(r.ndim): - k = k[:, None] + xyz = check_xyz_dim(xyz) + + r_vec = xyz - self.location + r = np.linalg.norm(r_vec, axis=-1, keepdims=True) + r_hat = r_vec/r + + k = append_ndim(self.wavenumber, r.ndim) kr = k * r ikr = 1j*kr @@ -378,7 +379,7 @@ def magnetic_field(self, xyz): self.current * self.length / (4 * np.pi * r**2) * (ikr + 1) * np.exp(-ikr) ) - return (-front_term[..., None] * np.cross(dxyz, self.orientation) / r[..., None]).squeeze() + return front_term * np.cross(self.orientation, r_hat) def magnetic_flux_density(self, xyz): r"""Magnetic flux density produced by the harmonic electric current dipole at a set of gridded locations. @@ -406,7 +407,7 @@ def magnetic_flux_density(self, xyz): Returns ------- - (n_freq, ..., 3) numpy.array of complex + (n_freq, ..., 3) numpy.ndarray of complex Magnetic flux at all frequencies for the gridded locations provided. Output array is squeezed when n_freq and/or n_loc = 1. @@ -473,11 +474,11 @@ def vector_potential(self, xyz): r"""Vector potential for the harmonic magnetic dipole at a set of gridded locations. For a harmonic magnetic dipole oriented in the :math:`\hat{u}` direction with - moment amplitude :math:`m`, the magnetic vector potential at frequency :math:`f` + moment amplitude :math:`m`, the electric vector potential at frequency :math:`f` at vector distance :math:`\mathbf{r}` from the dipole is given by: .. math:: - \mathbf{a}(\mathbf{r}) = \frac{i \omega \mu m}{4 \pi r} e^{-ikr} + \mathbf{F}(\mathbf{r}) = \frac{i \omega \mu m}{4 \pi r} e^{-ikr} \mathbf{\hat{u}} where @@ -487,15 +488,14 @@ def vector_potential(self, xyz): Parameters ---------- - xyz : (n, 3) numpy.ndarray + xyz : (..., 3) numpy.ndarray Gridded xyz locations Returns ------- - (n_freq, n_loc, 3) numpy.array of complex + (n_freq, ..., 3) numpy.ndarray of complex Magnetic vector potential at all frequencies for the gridded - locations provided. Output array is squeezed when n_freq and/or - n_loc = 1. + locations provided. Examples -------- @@ -544,14 +544,14 @@ def vector_potential(self, xyz): """ xyz = check_xyz_dim(xyz) - r = np.linalg.norm(xyz, axis=-1) - k = self.wavenumber - omega = self.omega - for i in range(r.ndim): - k = k[..., None] - omega = omega[..., None] + r_vec = xyz - self.location + r = np.linalg.norm(r_vec, axis=-1, keepdims=True) + + k = append_ndim(self.wavenumber, r.ndim) + omega = append_ndim(self.omega, r.ndim) + f = 1j * omega * self.mu * self.moment / (4 * np.pi * r) * np.exp(-1j * k * r) - return (f[..., None] * self.orientation).squeeze() + return f * self.orientation def electric_field(self, xyz): r"""Electric field for the harmonic magnetic dipole at a set of gridded locations. @@ -574,15 +574,14 @@ def electric_field(self, xyz): Parameters ---------- - xyz : (n, 3) numpy.ndarray + xyz : (..., 3) numpy.ndarray Gridded xyz locations Returns ------- - (n_freq, n_loc, 3) numpy.array of complex + (n_freq, ..., 3) numpy.ndarray of complex Electric field at all frequencies for the gridded - locations provided. Output array is squeezed when n_freq and/or - n_loc = 1. + locations provided. Examples -------- @@ -635,21 +634,20 @@ def electric_field(self, xyz): """ xyz = check_xyz_dim(xyz) - dxyz = xyz - self.location - r = np.linalg.norm(dxyz, axis=-1) - k = self.wavenumber - omega = self.omega - for i in range(r.ndim): - k = k[..., None] - omega = omega[..., None] + r_vec = xyz - self.location + r = np.linalg.norm(r_vec, axis=-1, keepdims=True) + r_hat = r_vec / r + k = append_ndim(self.wavenumber, r.ndim) + omega = append_ndim(self.omega, r.ndim) + kr = k * r ikr = 1j * kr front_term = ( (1j * omega * self.mu * self.moment) / (4. * np.pi * r**2) * (ikr + 1) * np.exp(-ikr) - ) / r - return (front_term[..., None] * np.cross(dxyz, self.orientation)).squeeze() + ) + return front_term * np.cross(r_hat, self.orientation) def current_density(self, xyz): r"""Current density for the harmonic magnetic dipole at a set of gridded locations. @@ -672,12 +670,12 @@ def current_density(self, xyz): Parameters ---------- - xyz : (n, 3) numpy.ndarray + xyz : (..., 3) numpy.ndarray Gridded xyz locations Returns ------- - (n_freq, n_loc, 3) numpy.array of complex + (n_freq, ..., 3) numpy.ndarray of complex Current density at all frequencies for the gridded locations provided. Output array is squeezed when n_freq and/or n_loc = 1. @@ -756,12 +754,12 @@ def magnetic_field(self, xyz): Parameters ---------- - xyz : (n, 3) numpy.ndarray + xyz : (..., 3) numpy.ndarray Gridded xyz locations Returns ------- - (n_freq, n_loc, 3) numpy.array of complex + (n_freq, ..., 3) numpy.ndarray of complex Magnetic field at all frequencies for the gridded locations provided. Output array is squeezed when n_freq and/or n_loc = 1. @@ -817,21 +815,21 @@ def magnetic_field(self, xyz): """ xyz = check_xyz_dim(xyz) - dxyz = xyz - self.location - r = np.linalg.norm(dxyz, axis=-1) - k = self.wavenumber - for i in range(r.ndim): - k = k[..., None] + r_vec = xyz - self.location + r = np.linalg.norm(r_vec, axis=-1, keepdims=True) + r_hat = r_vec / r + k = append_ndim(self.wavenumber, r.ndim) + kr = k * r ikr = 1j * kr front_term = self.moment / (4. * np.pi * r**3) * np.exp(-ikr) - symmetric_term = (dxyz @ self.orientation)[None, ...]*((-kr**2 + 3*ikr + 3) / r**2) + symmetric_term = (r_hat @ self.orientation)[..., None]*(-kr**2 + 3*ikr + 3) * r_hat - oriented_term = (kr**2 - ikr - 1)[..., None] * self.orientation + oriented_term = (kr**2 - ikr - 1) * self.orientation - return (front_term[..., None] * (symmetric_term[..., None] * dxyz + oriented_term)).squeeze() + return front_term * (symmetric_term + oriented_term) def magnetic_flux_density(self, xyz): r"""Magnetic flux density for the harmonic magnetic dipole at a set of gridded locations. @@ -855,12 +853,12 @@ def magnetic_flux_density(self, xyz): Parameters ---------- - xyz : (n, 3) numpy.ndarray + xyz : (..., 3) numpy.ndarray Gridded xyz locations Returns ------- - (n_freq, n_loc, 3) numpy.array of complex + (n_freq, ..., 3) numpy.ndarray of complex Magnetic flux density at all frequencies for the gridded locations provided. Output array is squeezed when n_freq and/or n_loc = 1. @@ -1017,7 +1015,7 @@ def electric_field(self, xyz): Returns ------- - (n_f, ..., 3) numpy.array of complex + (n_f, ..., 3) numpy.ndarray of complex Electric field at all frequencies for the gridded locations provided. @@ -1067,17 +1065,13 @@ def electric_field(self, xyz): >>> plt.show() """ xyz = check_xyz_dim(xyz) + z = xyz[..., [2]] - k = self.wavenumber + k = append_ndim(self.wavenumber, xyz.ndim) e0 = self.amplitude + ikz = 1j * k * z - z = xyz[..., 2] - for i in range(z.ndim): - k = k[..., None] - kz = k * z - ikz = 1j * kz - - return e0 * self.orientation * np.exp(ikz)[..., None] + return e0 * self.orientation * np.exp(ikz) def current_density(self, xyz): r"""Current density for the harmonic planewave at a set of gridded locations. @@ -1089,7 +1083,7 @@ def current_density(self, xyz): Returns ------- - (n_f, ..., 3) numpy.array of complex + (n_f, ..., 3) numpy.ndarray of complex Current density at all frequencies for the gridded locations provided. @@ -1158,7 +1152,7 @@ def magnetic_field(self, xyz): Returns ------- - (n_f, ..., 3) numpy.array of complex + (n_f, ..., 3) numpy.ndarray of complex Magnetic field at all frequencies for the gridded locations provided. @@ -1219,7 +1213,7 @@ def magnetic_flux_density(self, xyz): Returns ------- - (n_f, ..., 3) numpy.array of complex + (n_f, ..., 3) numpy.ndarray of complex Magnetic flux density at all frequencies for the gridded locations provided. @@ -1269,15 +1263,11 @@ def magnetic_flux_density(self, xyz): >>> plt.show() """ xyz = check_xyz_dim(xyz) + z = xyz[..., [2]] - k = self.wavenumber - omega = self.omega + k = append_ndim(self.wavenumber, xyz.ndim) + omega = append_ndim(self.omega, xyz.ndim) e0 = self.amplitude - - z = xyz[..., 2] - for i in range(z.ndim): - k = k[..., None] - omega = omega[..., None] kz = k * z ikz = 1j * kz @@ -1289,4 +1279,4 @@ def magnetic_flux_density(self, xyz): # account for the orientation in the cross product # take cross product with the propagation direction b_dir = np.cross(self.orientation, [0, 0, -1]) - return b_dir * b[..., None] + return b_dir * b diff --git a/geoana/em/static/__init__.py b/geoana/em/static/__init__.py index 87be330b..0014bf8c 100644 --- a/geoana/em/static/__init__.py +++ b/geoana/em/static/__init__.py @@ -14,6 +14,7 @@ :toctree: generated/ LineCurrentFreeSpace + LineCurrentWholeSpace MagneticDipoleWholeSpace MagneticPoleWholeSpace CircularLoopWholeSpace @@ -32,7 +33,8 @@ MagneticDipoleWholeSpace, MagneticPoleWholeSpace, CircularLoopWholeSpace, - PointCurrentWholeSpace + PointCurrentWholeSpace, + LineCurrentWholeSpace ) from geoana.em.static.halfspace import ( @@ -40,4 +42,8 @@ DipoleHalfSpace ) -from geoana.em.static.freespace import LineCurrentFreeSpace, MagneticPrism +from geoana.em.static.freespace import MagneticPrism + + +LineCurrentFreeSpace = LineCurrentWholeSpace + diff --git a/geoana/em/static/freespace.py b/geoana/em/static/freespace.py index 3c2c2405..9ed715cc 100644 --- a/geoana/em/static/freespace.py +++ b/geoana/em/static/freespace.py @@ -16,186 +16,6 @@ ) -class LineCurrentFreeSpace(BaseLineCurrent): - """Class for a static line current in free space. - - The ``LineCurrentFreeSpace`` class is used to analytically compute the - fields and potentials within freespace due to a set of static current-carrying wires. - """ - - def magnetic_field(self, xyz): - r"""Compute the magnetic field for the static current-carrying wire segments. - - Parameters - ---------- - xyz : (n, 3) numpy.ndarray xyz - gridded locations at which we are calculating the magnetic field - - Returns - ------- - (n, 3) numpy.ndarray - The magnetic field at each observation location in H/m. - - Examples - -------- - Here, we define a horizontal square loop and plot the magnetic field - on the xz-plane that intercepts at y=0. - - >>> from geoana.em.static import LineCurrentFreeSpace - >>> from geoana.utils import ndgrid - >>> from geoana.plotting_utils import plot2Ddata - >>> import numpy as np - >>> import matplotlib.pyplot as plt - - Let us begin by defining the loop. Note that to create an inductive - source, we closed the loop. - - >>> x_nodes = np.array([-0.5, 0.5, 0.5, -0.5, -0.5]) - >>> y_nodes = np.array([-0.5, -0.5, 0.5, 0.5, -0.5]) - >>> z_nodes = np.zeros_like(x_nodes) - >>> nodes = np.c_[x_nodes, y_nodes, z_nodes] - >>> simulation = LineCurrentFreeSpace(nodes) - - Now we create a set of gridded locations and compute the magnetic field. - - >>> xyz = ndgrid(np.linspace(-1, 1, 50), np.array([0]), np.linspace(-1, 1, 50)) - >>> H = simulation.magnetic_field(xyz) - - Finally, we plot the magnetic field. - - >>> fig = plt.figure(figsize=(4, 4)) - >>> ax = fig.add_axes([0.15, 0.15, 0.75, 0.75]) - >>> plot2Ddata(xyz[:, [0, 2]], H[:, [0, 2]], ax=ax, vec=True, scale='log', ncontour=25) - >>> ax.set_xlabel('X') - >>> ax.set_ylabel('Z') - >>> ax.set_title('Magnetic field') - - """ - - # TRANSMITTER NODES - I = self.current - tx_nodes = self.nodes - x1tr = tx_nodes[:, 0] - x2tr = tx_nodes[:, 1] - x3tr = tx_nodes[:, 2] - - n_loc = np.shape(xyz)[0] - n_segments = self.n_segments - - hx0 = np.zeros(n_loc) - hy0 = np.zeros(n_loc) - hz0 = np.zeros(n_loc) - - for pp in range(0, n_segments): - - # Wire ends for transmitter wire pp - x1a = x1tr[pp] - x2a = x2tr[pp] - x3a = x3tr[pp] - x1b = x1tr[pp + 1] - x2b = x2tr[pp + 1] - x3b = x3tr[pp + 1] - - # Vector Lengths between points - vab = np.sqrt((x1b - x1a) ** 2 + (x2b - x2a) ** 2 + (x3b - x3a) ** 2) - vap = np.sqrt( - (xyz[:, 0] - x1a) ** 2 + (xyz[:, 1] - x2a) ** 2 + (xyz[:, 2] - x3a) ** 2 - ) - vbp = np.sqrt( - (xyz[:, 0] - x1b) ** 2 + (xyz[:, 1] - x2b) ** 2 + (xyz[:, 2] - x3b) ** 2 - ) - - # Cosines from cos()=/(|v1||v2|) - cos_alpha = ( - (xyz[:, 0] - x1a) * (x1b - x1a) - + (xyz[:, 1] - x2a) * (x2b - x2a) - + (xyz[:, 2] - x3a) * (x3b - x3a) - ) / (vap * vab) - cos_beta = ( - (xyz[:, 0] - x1b) * (x1a - x1b) - + (xyz[:, 1] - x2b) * (x2a - x2b) - + (xyz[:, 2] - x3b) * (x3a - x3b) - ) / (vbp * vab) - - # Determining Radial Vector From Wire - dot_temp = ( - (x1a - xyz[:, 0]) * (x1b - x1a) - + (x2a - xyz[:, 1]) * (x2b - x2a) - + (x3a - xyz[:, 2]) * (x3b - x3a) - ) - - rx1 = (x1a - xyz[:, 0]) - dot_temp * (x1b - x1a) / vab ** 2 - rx2 = (x2a - xyz[:, 1]) - dot_temp * (x2b - x2a) / vab ** 2 - rx3 = (x3a - xyz[:, 2]) - dot_temp * (x3b - x3a) / vab ** 2 - - r = np.sqrt(rx1 ** 2 + rx2 ** 2 + rx3 ** 2) - - phi = (cos_alpha + cos_beta) / r - - # I/4*pi in each direction - ix1 = I * (x1b - x1a) / (4 * np.pi * vab) - ix2 = I * (x2b - x2a) / (4 * np.pi * vab) - ix3 = I * (x3b - x3a) / (4 * np.pi * vab) - - # Add contribution from wire pp into array - hx0 = hx0 + phi * (-ix2 * rx3 + ix3 * rx2) / r - hy0 = hy0 + phi * (ix1 * rx3 - ix3 * rx1) / r - hz0 = hz0 + phi * (-ix1 * rx2 + ix2 * rx1) / r - - return np.c_[hx0, hy0, hz0] - - def magnetic_flux_density(self, xyz): - r"""Compute the magnetic flux density for the static current-carrying wire segments. - - Parameters - ---------- - xyz : (n, 3) numpy.ndarray xyz - gridded locations at which we are calculating the magnetic flux density - - Returns - ------- - (n, 3) numpy.ndarray - The magnetic flux density at each observation location in T. - - Examples - -------- - Here, we define a horizontal square loop and plot the magnetic flux - density on the XZ-plane that intercepts at Y=0. - - >>> from geoana.em.static import LineCurrentFreeSpace - >>> from geoana.utils import ndgrid - >>> from geoana.plotting_utils import plot2Ddata - >>> import numpy as np - >>> import matplotlib.pyplot as plt - - Let us begin by defining the loop. Note that to create an inductive - source, we closed the loop - - >>> x_nodes = np.array([-0.5, 0.5, 0.5, -0.5, -0.5]) - >>> y_nodes = np.array([-0.5, -0.5, 0.5, 0.5, -0.5]) - >>> z_nodes = np.zeros_like(x_nodes) - >>> nodes = np.c_[x_nodes, y_nodes, z_nodes] - >>> simulation = LineCurrentFreeSpace(nodes) - - Now we create a set of gridded locations and compute the magnetic flux density. - - >>> xyz = ndgrid(np.linspace(-1, 1, 50), np.array([0]), np.linspace(-1, 1, 50)) - >>> B = simulation.magnetic_flux_density(xyz) - - Finally, we plot the magnetic flux density on the plane. - - >>> fig = plt.figure(figsize=(4, 4)) - >>> ax = fig.add_axes([0.15, 0.15, 0.8, 0.8]) - >>> plot2Ddata(xyz[:, [0, 2]], B[:, [0, 2]], ax=ax, vec=True, scale='log', ncontour=25) - >>> ax.set_xlabel('X') - >>> ax.set_ylabel('Z') - >>> ax.set_title('Magnetic flux density') - - """ - - return mu_0 * self.magnetic_field(xyz) - - class MagneticPrism(BasePrism): """Class for magnetic field solutions for a prism. diff --git a/geoana/em/static/wholespace.py b/geoana/em/static/wholespace.py index d27a8106..facccb1e 100644 --- a/geoana/em/static/wholespace.py +++ b/geoana/em/static/wholespace.py @@ -1,7 +1,8 @@ import numpy as np from scipy.special import ellipk, ellipe +from scipy.spatial.transform import Rotation -from ..base import BaseDipole, BaseMagneticDipole, BaseEM +from ..base import BaseDipole, BaseMagneticDipole, BaseEM, BaseLineCurrent from ... import spatial from geoana.utils import check_xyz_dim @@ -10,6 +11,10 @@ "MagneticPoleWholeSpace", "PointCurrentWholeSpace" ] +from ...kernels import prism_fzx, prism_fzy +from ...spatial import cartesian_2_cylindrical, cylindrical_2_cartesian, cylindrical_to_cartesian, \ + cartesian_2_spherical, spherical_to_cartesian, cartesian_to_spherical, cartesian_to_cylindrical + class MagneticDipoleWholeSpace(BaseEM, BaseMagneticDipole): """Class for a static magnetic dipole in a wholespace. @@ -43,7 +48,7 @@ def vector_potential(self, xyz, coordinates="cartesian"): Parameters ---------- - xyz : (n, 3) numpy.ndarray xyz + xyz : (..., 3) numpy.ndarray xyz gridded locations at which we are calculating the vector potential coordinates: str {'cartesian', 'cylindrical'} coordinate system that the location (xyz) are provided. @@ -52,7 +57,7 @@ def vector_potential(self, xyz, coordinates="cartesian"): Returns ------- - (n, 3) numpy.ndarray + (..., 3) numpy.ndarray The magnetic vector potential at each observation location in the coordinate system specified in units *Tm*. @@ -101,18 +106,15 @@ def vector_potential(self, xyz, coordinates="cartesian"): ) xyz = check_xyz_dim(xyz) - n_obs = xyz.shape[0] - # orientation of the dipole if coordinates.lower() == "cylindrical": xyz = spatial.cylindrical_2_cartesian(xyz) - dxyz = self.vector_distance(xyz) - r = spatial.repeat_scalar(self.distance(xyz)) - m = self.moment * np.atleast_2d(self.orientation).repeat(n_obs, axis=0) + r_vec = xyz - self.location + r = np.linalg.norm(r_vec, axis=-1, keepdims=True) + r_hat = r_vec / r - m_cross_r = np.cross(m, dxyz) - a = (self.mu / (4 * np.pi)) * m_cross_r / (r**3) + a = self.moment * self.mu / (4 * np.pi * r**2) * np.cross(self.orientation, r_hat) if coordinates.lower() == "cylindrical": a = spatial.cartesian_2_cylindrical(xyz, a) @@ -143,7 +145,7 @@ def magnetic_flux_density(self, xyz, coordinates="cartesian"): Parameters ---------- - xyz : (n, 3) numpy.ndarray + xyz : (..., 3) numpy.ndarray gridded locations at which we calculate the magnetic flux density coordinates: str {'cartesian', 'cylindrical'} coordinate system that the location (xyz) are provided. @@ -152,7 +154,7 @@ def magnetic_flux_density(self, xyz, coordinates="cartesian"): Returns ------- - (n, 3) numpy.ndarray + (..., 3) numpy.ndarray The magnetic flux density at each observation location in the coordinate system specified in Teslas. @@ -206,14 +208,11 @@ def magnetic_flux_density(self, xyz, coordinates="cartesian"): xyz = spatial.cylindrical_2_cartesian(xyz) r_vec = xyz - self.location - r = np.linalg.norm(r_vec, axis=-1)[..., None] - m_vec = self.moment * self.orientation - - m_dot_r = np.einsum('...i,i->...', r_vec, m_vec)[..., None] + r = np.linalg.norm(r_vec, axis=-1, keepdims=True) + r_hat = r_vec / r - b = (self.mu / (4 * np.pi)) * ( - (3.0 * r_vec * m_dot_r / (r ** 5)) - - m_vec / (r ** 3) + b = self.moment * self.mu / (4 * np.pi * r**3) * ( + 3 * r_hat.dot(self.orientation)[..., None] * r_hat - self.orientation ) if coordinates.lower() == "cylindrical": @@ -244,7 +243,7 @@ def magnetic_field(self, xyz, coordinates="cartesian"): Parameters ---------- - xyz : (n, 3) numpy.ndarray xyz + xyz : (..., 3) numpy.ndarray xyz gridded locations at which we calculate the magnetic field coordinates: str {'cartesian', 'cylindrical'} coordinate system that the location (xyz) are provided. @@ -253,7 +252,7 @@ def magnetic_field(self, xyz, coordinates="cartesian"): Returns ------- - (n, 3) numpy.ndarray + (..., 3) numpy.ndarray The magnetic field at each observation location in the coordinate system specified in units A/m. @@ -325,7 +324,7 @@ def magnetic_flux_density(self, xyz, coordinates="cartesian"): Parameters ---------- - xyz : (n, 3) numpy.ndarray xyz + xyz : (..., 3) numpy.ndarray xyz gridded xyz locations at which we calculate the magnetic flux density coordinates: str {'cartesian', 'cylindrical'} coordinate system that the location (xyz) are provided. @@ -334,7 +333,7 @@ def magnetic_flux_density(self, xyz, coordinates="cartesian"): Returns ------- - (n, 3) numpy.ndarray + (..., 3) numpy.ndarray The magnetic flux density at each observation location in the coordinate system specified in units T. @@ -352,10 +351,11 @@ def magnetic_flux_density(self, xyz, coordinates="cartesian"): if coordinates.lower() == "cylindrical": xyz = spatial.cylindrical_2_cartesian(xyz) - r = self.vector_distance(xyz) - dxyz = spatial.repeat_scalar(self.distance(xyz)) + r_vec = xyz - self.location + r = np.linalg.norm(r_vec, axis=-1, keepdims=True) + r_hat = r_vec / r - b = self.moment * self.mu / (4 * np.pi * (dxyz ** 3)) * r + b = self.moment * self.mu / (4 * np.pi * r ** 2) * r_hat if coordinates.lower() == "cylindrical": b = spatial.cartesian_2_cylindrical(xyz, b) @@ -384,7 +384,7 @@ def magnetic_field(self, xyz, coordinates="cartesian"): Parameters ---------- - xyz : (n, 3) numpy.ndarray xyz + xyz : (..., 3) numpy.ndarray xyz gridded locations at which we calculate the magnetic field coordinates: str {'cartesian', 'cylindrical'} coordinate system that the location (xyz) are provided. @@ -393,7 +393,7 @@ def magnetic_field(self, xyz, coordinates="cartesian"): Returns ------- - (n, 3) numpy.ndarray + (..., 3) numpy.ndarray The magnetic field at each observation location in the coordinate system specified in units A/m. @@ -493,7 +493,7 @@ def vector_potential(self, xyz, coordinates="cartesian"): .. math:: a_\theta (\rho, z) = \frac{\mu_0 I}{\pi k} - \sqrt{ \frac{R}{\rho^2}} \bigg [ (1 - k^2/2) \, K(k^2) - K(k^2) \bigg ] + \sqrt{ \frac{R}{\rho^2}} \bigg [ (1 - k^2/2) \, K(k^2) - E(k^2) \bigg ] where @@ -510,7 +510,7 @@ def vector_potential(self, xyz, coordinates="cartesian"): Parameters ---------- - xyz : (n, 3) numpy.ndarray xyz + xyz : (..., 3) numpy.ndarray xyz gridded locations at which we calculate the vector potential coordinates: str {'cartesian', 'cylindrical'} coordinate system that the location (xyz) are provided. @@ -519,7 +519,7 @@ def vector_potential(self, xyz, coordinates="cartesian"): Returns ------- - (n, 3) numpy.ndarray + (..., 3) numpy.ndarray The magnetic vector potential at each observation location in the coordinate system specified in units *Tm*. @@ -574,43 +574,62 @@ def vector_potential(self, xyz, coordinates="cartesian"): if coordinates.lower() == "cylindrical": xyz = spatial.cylindrical_2_cartesian(xyz) - xyz = spatial.rotate_points_from_normals( - xyz, np.array(self.orientation), # work around for a properties issue - np.r_[0., 0., 1.], x0=np.array(self.location) - ) + # define a rotation matrix that rotates my orientation to z: + rot, _ = Rotation.align_vectors(np.array([0, 0, 1]), self.orientation) - n_obs = xyz.shape[0] - dxyz = self.vector_distance(xyz) - r = self.distance(xyz) + # rotate the points + r_vec = rot.apply(xyz.reshape(-1, 3) - self.location) - rho = np.sqrt((dxyz[:, :2]**2).sum(1)) + r_cyl = cartesian_2_cylindrical(r_vec) - k2 = (4 * self.radius * rho) / ((self.radius + rho)**2 +dxyz[:, 2]**2) - k2[k2 > 1.] = 1. # if there are any rounding errors + rho = r_cyl[..., 0] + z = r_cyl[..., 2] - E = ellipe(k2) - K = ellipk(k2) + a = self.radius + C = self.mu * self.current / np.pi - # singular if rho = 0, k2 = 1 - ind = (rho > eps) & (k2 < 1) + alpha_sq = (a - rho)**2 + z**2 + beta_sq = (a + rho)**2 + z**2 + beta = np.sqrt(beta_sq) - Atheta = np.zeros_like(r) - Atheta[ind] = ( - (self.mu * self.current) / (np.pi * np.sqrt(k2[ind])) * - np.sqrt(self.radius / rho[ind]) * - ((1. - k2[ind] / 2.)*K[ind] - E[ind]) - ) + k_sq = 1 - alpha_sq/beta_sq + + ek = ellipe(k_sq) + kk = ellipk(k_sq) + + A_cyl = np.zeros_like(r_vec) - # assume that the z-axis aligns with the polar axis - A = np.zeros_like(xyz) - A[ind, 0] = Atheta[ind] * (-dxyz[ind, 1] / rho[ind]) - A[ind, 1] = Atheta[ind] * (dxyz[ind, 0] / rho[ind]) + # when rho is small relative to the radius and z, k_sq -> 0 + # this is unstable so use a small argument approximation. + # check k_sq for the relative sizes + small_rho = k_sq < 1E-3 - # rotate the points to aligned with the normal to the source - A = spatial.rotate_points_from_normals( - A, np.r_[0., 0., 1.], np.array(self.orientation), - x0=np.array(self.location) + temp = np.sqrt(a**2 + z[small_rho]**2) + A_cyl[small_rho, 1] = np.pi * C * ( + # A(rho=0) = 0 + rho[small_rho] * a**2 / (4 * temp**3) + # rho * A'(rho=0) + # A''(rho=0) = 0 + rho[small_rho]**3 * 3 * a**2 / ( a ** 2 - 4 * z[small_rho]**2) / ( + 32 * np.sqrt(a**2 + z[small_rho]**2)**7 + ) ) + z = z[~small_rho] + beta = beta[~small_rho] + beta_sq = beta_sq[~small_rho] + rho = rho[~small_rho] + ek = ek[~small_rho] + kk = kk[~small_rho] + + # singular if alpha = 0 + # (a.k.a. the point was numerically on the loop) + A_cyl[~small_rho, 1] = C / (2 * rho * beta) * ( + (a**2 + rho**2 + z**2) * kk - beta_sq * ek + ) + + A = cylindrical_to_cartesian(r_cyl, A_cyl) + + # un-do the rotation on the vector components. + A = rot.apply(A, inverse=True).reshape(xyz.shape) if coordinates.lower() == "cylindrical": A = spatial.cartesian_2_cylindrical(xyz, A) @@ -682,6 +701,7 @@ def magnetic_flux_density(self, xyz, coordinates="cartesian"): """ xyz = check_xyz_dim(xyz) + # convert coordinates if not cartesian if coordinates.lower() == "cylindrical": xyz = spatial.cylindrical_2_cartesian(xyz) @@ -691,55 +711,62 @@ def magnetic_flux_density(self, xyz, coordinates="cartesian"): f"system you provided, {coordinates}, is not yet supported." ) - xyz = spatial.rotate_points_from_normals( - xyz, np.array(self.orientation), # work around for a properties issue - np.r_[0., 0., 1.], x0=np.array(self.location) - ) - # rotate all the points such that the orientation is directly vertical - - dxyz = self.vector_distance(xyz) - r = self.distance(xyz) + # define a rotation matrix that rotates my orientation to z: + rot, _ = Rotation.align_vectors(np.array([0, 0, 1]), self.orientation) - rho = np.linalg.norm(dxyz[:, :2], axis=-1) + # rotate the points + r_vec = rot.apply(xyz.reshape(-1, 3) - self.location) - B = np.zeros((len(rho), 3)) + r_cyl = cartesian_to_cylindrical(r_vec) + rho = r_cyl[:, 0] + z = r_cyl[:, 2] + a = self.radius + C = self.mu * self.current / np.pi - # for On axis points - inds_axial = rho==0.0 + alpha_sq = (a - rho)**2 + z**2 + beta_sq = (a + rho)**2 + z**2 + beta = np.sqrt(beta_sq) + k_sq = 1 - alpha_sq / beta_sq - B[inds_axial, -1] = self.mu * self.current * self.radius**2 / ( - 2 * (self.radius**2 + dxyz[inds_axial, 2]**2)**(1.5) - ) + ek = ellipe(k_sq) + kk = ellipk(k_sq) - # Off axis - alpha = rho[~inds_axial]/self.radius - beta = dxyz[~inds_axial, 2]/self.radius - gamma = dxyz[~inds_axial, 2]/rho[~inds_axial] + B_cyl = np.zeros_like(r_cyl) - Q = ((1+alpha)**2 + beta**2) - k2 = 4 * alpha/Q + # when rho is small relative to the radius and z, k_sq -> 0 + # this is unstable so use a small argument approximation. + # check k_sq for the relative sizes + small_rho = k_sq < 1E-3 - # axial part: - B[~inds_axial, -1] = self.mu * self.current / (2 * self.radius * np.pi * np.sqrt(Q)) * ( - ellipe(k2)*(1 - alpha**2 - beta**2)/(Q - 4 * alpha) + ellipk(k2) + temp = np.sqrt(a**2 + z[small_rho]**2) + B_cyl[small_rho, 0] = 3 * C * np.pi * a**2 * z[small_rho] * ( + rho[small_rho] /(4 * temp ** 5) + + rho[small_rho]**3 * (15 * a**2 - 20 * z[small_rho]**2)/(32 * temp**9) ) - # radial part: - B_rad = self.mu * self.current * gamma / (2 * self.radius * np.pi * np.sqrt(Q)) * ( - ellipe(k2)*(1 + alpha**2 + beta**2)/(Q - 4 * alpha) - ellipk(k2) + # this expectedly blows up when rho = radius and z = 0 + # Meaning it is literally on the loop... + B_cyl[:, 2] = C / (2 * alpha_sq * beta) * ( + (a**2 - rho**2 - z**2) * ek + alpha_sq * kk ) - # convert radial component to x and y.. - B[~inds_axial, 0] = B_rad * (dxyz[~inds_axial, 0]/rho[~inds_axial]) - B[~inds_axial, 1] = B_rad * (dxyz[~inds_axial, 1]/rho[~inds_axial]) + z = z[~small_rho] + alpha_sq = alpha_sq[~small_rho] + beta = beta[~small_rho] + rho = rho[~small_rho] + ek = ek[~small_rho] + kk = kk[~small_rho] - # rotate the vectors to be aligned with the normal to the source - B = spatial.rotate_points_from_normals( - B, np.r_[0., 0., 1.], np.array(self.orientation), + B_cyl[~small_rho, 0] = C * z / (2 * alpha_sq * beta * rho) * ( + (a**2 + rho**2 + z**2) * ek - alpha_sq * kk ) + B = cylindrical_to_cartesian(r_cyl, B_cyl) + + B = rot.apply(B, inverse=True).reshape(xyz.shape) + if coordinates.lower() == "cylindrical": - B = spatial.cartesian_2_cylindrical(xyz, B) + B = cartesian_to_cylindrical(xyz, B) return B @@ -810,74 +837,272 @@ def magnetic_field(self, xyz, coordinates="cartesian"): return self.magnetic_flux_density(xyz, coordinates=coordinates) / self.mu -class PointCurrentWholeSpace: - """Class for a point current in a wholespace. - - The ``PointCurrentWholeSpace`` class is used to analytically compute the - potentials, current densities and electric fields within a wholespace due to a point current. +class LineCurrentWholeSpace(BaseLineCurrent, BaseEM): + """Class for a static line current in whole space. - Parameters - ---------- - current : float - Electrical current in the point current (A). Default is 1A. - rho : float - Resistivity in the point current (:math:`\\Omega \\cdot m`). - location : array_like, optional - Location at which we are observing in 3D space (m). Default is (0, 0, 0). + The ``LineCurrentFreeSpace`` class is used to analytically compute the + fields and potentials within a wholespace due to a set of constant current-carrying + wires. """ - def __init__(self, rho, current=1.0, location=None): + def scalar_potential(self, xyz): + xyz = check_xyz_dim(xyz) + # If i had a single point, treat it as a source + if self.n_segments == 0: + r = np.linalg.norm(xyz - self.nodes[0], axis=-1) + return self.current/(4 * np.pi * self.sigma * r) + # if I was a closed loop, return 0 + if np.all(self.nodes[-1] == self.nodes[0]): + return np.zeros_like(xyz[..., 0]) - self.current = current - self.rho = rho - if location is None: - location = np.r_[0, 0, 0] - self.location = location + r_A = np.linalg.norm(xyz - self.nodes[-1], axis=-1) + r_B = np.linalg.norm(xyz - self.nodes[0], axis=-1) - @property - def current(self): - """Current in the point current in Amps. + return self.current/(4 * np.pi * self.sigma) * (1/r_A - 1/r_B) + + def vector_potential(self, xyz): + + xyz = check_xyz_dim(xyz) + + out = np.zeros_like(xyz) + temp_storage = np.zeros_like(xyz) + for p0, p1 in zip(self.nodes[:-1], self.nodes[1:]): + l_vec = p1 - p0 + l = np.linalg.norm(l_vec) + l_hat = l_vec / l + + # find the rotation from the line segments orientation + # to the x_hat direction. + rot, _ = Rotation.align_vectors([1, 0, 0], l_hat.reshape(-1, 3)) + + # shift and rotate the grid points + r0_vec = rot.apply(xyz.reshape(-1, 3) - p0).reshape(xyz.shape) + + #p1 would've been shifted and rotated to [l, 0, 0] + r1_vec = r0_vec - np.array([l, 0, 0]) + + # hey these are just the 1/r kernel evaluations! + v0_x = -prism_fzy(r0_vec[..., 0], r0_vec[..., 1], r0_vec[..., 2]) + v1_x = -prism_fzy(r1_vec[..., 0], r1_vec[..., 1], r1_vec[..., 2]) + + temp_storage[..., 0] = v1_x - v0_x + # the undo the local rotation... + out += rot.apply(temp_storage.reshape(-1, 3), inverse=True).reshape(xyz.shape) + + out *= -self.mu * self.current / (4 * np.pi) + + # note because this is a whole space, we do not have to deal with the + # magnetic fields due to the current flowing out of the ends of a grounded + # wire. + return out + + def magnetic_field(self, xyz): + r"""Compute the magnetic field for the static current-carrying wire segments. + + Parameters + ---------- + xyz : (..., 3) numpy.ndarray xyz + gridded locations at which we are calculating the magnetic field Returns ------- - float - Current in the point current in Amps. + (..., 3) numpy.ndarray + The magnetic field at each observation location in H/m. + + Examples + -------- + Here, we define a horizontal square loop and plot the magnetic field + on the xz-plane that intercepts at y=0. + + >>> from geoana.em.static import LineCurrentWholeSpace + >>> from geoana.utils import ndgrid + >>> from geoana.plotting_utils import plot2Ddata + >>> import numpy as np + >>> import matplotlib.pyplot as plt + + Let us begin by defining the loop. Note that to create an inductive + source, we closed the loop. + + >>> x_nodes = np.array([-0.5, 0.5, 0.5, -0.5, -0.5]) + >>> y_nodes = np.array([-0.5, -0.5, 0.5, 0.5, -0.5]) + >>> z_nodes = np.zeros_like(x_nodes) + >>> nodes = np.c_[x_nodes, y_nodes, z_nodes] + >>> simulation = LineCurrentWholeSpace(nodes) + + Now we create a set of gridded locations and compute the magnetic field. + + >>> xyz = ndgrid(np.linspace(-1, 1, 50), np.array([0]), np.linspace(-1, 1, 50)) + >>> H = simulation.magnetic_field(xyz) + + Finally, we plot the magnetic field. + + >>> fig = plt.figure(figsize=(4, 4)) + >>> ax = fig.add_axes([0.15, 0.15, 0.75, 0.75]) + >>> plot2Ddata(xyz[:, [0, 2]], H[:, [0, 2]], ax=ax, vec=True, scale='log', ncontour=25) + >>> ax.set_xlabel('X') + >>> ax.set_ylabel('Z') + >>> ax.set_title('Magnetic field') + """ - return self._current + return self.magnetic_flux_density(xyz) / self.mu - @current.setter - def current(self, value): + def magnetic_flux_density(self, xyz): + r"""Compute the magnetic flux density for the static current-carrying wire segments. - try: - value = float(value) - except: - raise TypeError(f"current must be a number, got {type(value)}") + Parameters + ---------- + xyz : (..., 3) numpy.ndarray xyz + gridded locations at which we are calculating the magnetic flux density - self._current = value + Returns + ------- + (..., 3) numpy.ndarray + The magnetic flux density at each observation location in T. - @property - def rho(self): - """Resistivity in the point current in :math:`\\Omega \\cdot m` + Examples + -------- + Here, we define a horizontal square loop and plot the magnetic flux + density on the XZ-plane that intercepts at Y=0. + + >>> from geoana.em.static import LineCurrentWholeSpace + >>> from geoana.utils import ndgrid + >>> from geoana.plotting_utils import plot2Ddata + >>> import numpy as np + >>> import matplotlib.pyplot as plt + + Let us begin by defining the loop. Note that to create an inductive + source, we closed the loop + + >>> x_nodes = np.array([-0.5, 0.5, 0.5, -0.5, -0.5]) + >>> y_nodes = np.array([-0.5, -0.5, 0.5, 0.5, -0.5]) + >>> z_nodes = np.zeros_like(x_nodes) + >>> nodes = np.c_[x_nodes, y_nodes, z_nodes] + >>> simulation = LineCurrentWholeSpace(nodes) + + Now we create a set of gridded locations and compute the magnetic flux density. + + >>> xyz = ndgrid(np.linspace(-1, 1, 50), np.array([0]), np.linspace(-1, 1, 50)) + >>> B = simulation.magnetic_flux_density(xyz) + + Finally, we plot the magnetic flux density on the plane. + + >>> fig = plt.figure(figsize=(4, 4)) + >>> ax = fig.add_axes([0.15, 0.15, 0.8, 0.8]) + >>> plot2Ddata(xyz[:, [0, 2]], B[:, [0, 2]], ax=ax, vec=True, scale='log', ncontour=25) + >>> ax.set_xlabel('X') + >>> ax.set_ylabel('Z') + >>> ax.set_title('Magnetic flux density') + + """ + + xyz = check_xyz_dim(xyz) + + out = np.zeros_like(xyz) + temp_storage = np.zeros_like(xyz) + for p0, p1 in zip(self.nodes[:-1], self.nodes[1:]): + l_vec = p1 - p0 + l = np.linalg.norm(l_vec) + l_hat = l_vec / l + + # find the rotation from the line segments orientation + # to the x_hat direction. + rot, _ = Rotation.align_vectors([1, 0, 0], l_hat) + + # shift and rotate the grid points + r0_vec = rot.apply(xyz.reshape(-1, 3) - p0).reshape(xyz.shape) + #p1 would've been shifted and rotated to [l, 0, 0] + r1_vec = r0_vec - np.array([l, 0, 0]) + + r0 = np.linalg.norm(r0_vec, axis=-1, keepdims=True) + r1 = np.linalg.norm(r1_vec, axis=-1, keepdims=True) + r0_hat = r0_vec / r0 + r1_hat = r1_vec / r1 + + cyl_points = cartesian_2_cylindrical(r0_vec[..., 1:]) + + temp_storage[..., 1] = (r1_hat[..., 0] - r0_hat[..., 0])/cyl_points[..., 0] + temp_storage[..., 1:] = cylindrical_2_cartesian(cyl_points, temp_storage) + + # the undo the local rotation... + out += rot.apply(temp_storage.reshape(-1, 3), inverse=True).reshape(xyz.shape) + + out *= -self.mu * self.current / (4 * np.pi) + return out + + def electric_field(self, xyz): + r"""Compute the electric for the static current-carrying wire segments. + + If the wire is closed, there is no electric field, but if it is an open + wire, + + Parameters + ---------- + xyz : (..., 3) numpy.ndarray xyz + gridded locations at which we are calculating the electric field Returns ------- - float - Resistivity in the point current in :math:`\\Omega \\cdot m` + (..., 3) numpy.ndarray + The electric field y at each observation location in T. + """ - return self._rho + xyz = check_xyz_dim(xyz) + # If I had a single point, treat it as a source + if self.n_segments == 0: + r_vec = xyz - self.nodes[0] + r = np.linalg.norm(r_vec, axis=-1, keepdims=True) + return self.current / (4 * np.pi * self.sigma) * (r_vec/r**3) + # if I was a closed loop, return 0 + if np.all(self.nodes[-1] == self.nodes[0]): + return np.zeros_like(xyz) + # otherwise, current leaks out at the ends! + r_vec_A = xyz - self.nodes[-1] + r_vec_B = xyz - self.nodes[0] + r_A = np.linalg.norm(r_vec_A, axis=-1, keepdims=True) + r_B = np.linalg.norm(r_vec_B, axis=-1, keepdims=True) + + return self.current/(4 * np.pi * self.sigma) * (r_vec_A/r_A**3 - r_vec_B/r_B**3) - @rho.setter - def rho(self, value): + def current_density(self, xyz): + r"""Compute the current density for the static current-carrying wire segments. - try: - value = float(value) - except: - raise TypeError(f"current must be a number, got {type(value)}") + If the wire is closed, there is no current density, but if it is an open + wire, - if value <= 0.0: - raise ValueError("current must be greater than 0") + Parameters + ---------- + xyz : (..., 3) numpy.ndarray xyz + gridded locations at which we are calculating the current density + + Returns + ------- + (..., 3) numpy.ndarray + The current density y at each observation location in T. + + """ + return self.sigma * self.electric_field(xyz) + + +class PointCurrentWholeSpace(LineCurrentWholeSpace): + """Class for a point current in a wholespace. + + The ``PointCurrentWholeSpace`` class is used to analytically compute the + potentials, current densities and electric fields within a wholespace due to a point current. - self._rho = value + Parameters + ---------- + rho : float + Resistivity in the point current (:math:`\\Omega \\cdot m`). + current : float + Electrical current in the point current (A). Default is 1A. + location : array_like, optional + Location at which we are observing in 3D space (m). Default is (0, 0, 0). + """ + + def __init__(self, rho, current=1.0, location=None): + if location is None: + location = np.r_[0, 0, 0] + super().__init__(sigma=1/rho, nodes=location, current=current) @property def location(self): @@ -888,7 +1113,7 @@ def location(self): (3) numpy.ndarray of float Location of observer in 3D space. Default = np.r_[0,0,0]. """ - return self._location + return self.nodes[0] @location.setter def location(self, vec): @@ -902,8 +1127,7 @@ def location(self, vec): raise ValueError( f"location must be array_like with shape (3,), got {len(vec)}" ) - - self._location = vec + self.nodes[0] = vec def potential(self, xyz): """Electric potential for a point current in a wholespace. @@ -962,13 +1186,7 @@ def potential(self, xyz): >>> plt.title('Electric Potential from Point Current in a Wholespace') >>> plt.show() """ - - xyz = check_xyz_dim(xyz) - r_vec = xyz - self.location - r = np.linalg.norm(r_vec, axis=-1) - - v = self.rho * self.current / (4 * np.pi * r) - return v + return super().scalar_potential(xyz) def electric_field(self, xyz): """Electric field for a point current in a wholespace. @@ -1028,13 +1246,7 @@ def electric_field(self, xyz): >>> plt.title('Electric Field Lines for a Point Current in a Wholespace') >>> plt.show() """ - - xyz = check_xyz_dim(xyz) - r_vec = xyz - self.location - r = np.linalg.norm(r_vec, axis=-1) - - e = self.rho * self.current * r_vec / (4 * np.pi * r[..., None] ** 3) - return e + return super().electric_field(xyz) def current_density(self, xyz): """Current density for a point current in a wholespace. @@ -1093,6 +1305,4 @@ def current_density(self, xyz): >>> plt.title('Current Density for a Point Current in a Wholespace') >>> plt.show() """ - - j = self.electric_field(xyz) / self.rho - return j + return super().current_density(xyz) diff --git a/geoana/em/tdem/base.py b/geoana/em/tdem/base.py index 36fe320f..ccc97c63 100644 --- a/geoana/em/tdem/base.py +++ b/geoana/em/tdem/base.py @@ -159,7 +159,7 @@ def time(self, value): # Ensure float or numpy array of float try: - value = np.atleast_1d(value).astype(float) + value = np.asarray(value, dtype=float) except: raise TypeError(f"times are not a valid type") @@ -167,7 +167,7 @@ def time(self, value): if (value < 0.).any(): raise ValueError("All times must be greater than 0") if value.ndim > 1: - raise TypeError(f"times must be ('*') array") + raise TypeError(f"times must have at most 1 dimension.") self._time = value diff --git a/geoana/em/tdem/halfspace.py b/geoana/em/tdem/halfspace.py index 34873c05..ae09fc3d 100644 --- a/geoana/em/tdem/halfspace.py +++ b/geoana/em/tdem/halfspace.py @@ -4,6 +4,7 @@ from geoana.em.tdem.base import BaseTDEM from geoana.em.tdem.simple_functions import magnetic_field_vertical_magnetic_dipole, magnetic_field_time_deriv_magnetic_dipole +from geoana.utils import check_xyz_dim class VerticalMagneticDipoleHalfSpace(BaseTDEM, BaseMagneticDipole): @@ -23,19 +24,23 @@ def magnetic_field(self, xy): Parameters ---------- - xy : (n_locations, 2) numpy.ndarray + xy : (..., 2) numpy.ndarray receiver locations Returns ------- - numpy.ndarray + (n_t, ..., 3) numpy.ndarray magnetic field for each xy location """ - dxy = xy - self.location + try: + xy = check_xyz_dim(xy, dim=2) + except ValueError: + xy = check_xyz_dim(xy, dim=3)[..., :2] + dxy = xy - self.location[:2] h = magnetic_field_vertical_magnetic_dipole( - np.r_[self.time], dxy[:, :2], self.sigma, self.mu, self.moment + self.time, dxy, self.sigma, self.mu, self.moment ) - return h[0] # because time was a 1 element array + return h def magnetic_flux_density(self, xy): """Magnetic flux due to a step off magnetic dipole over a half space @@ -45,12 +50,12 @@ def magnetic_flux_density(self, xy): Parameters ---------- - xy : (n_locations, 2) numpy.ndarray + xy : (..., 2) numpy.ndarray receiver locations Returns ------- - numpy.ndarray + (n_t, ..., 3) numpy.ndarray magnetic flux for each xy location """ return self.mu * self.magnetic_field(xy) @@ -63,20 +68,23 @@ def magnetic_field_time_derivative(self, xy): Parameters ---------- - xy : (n_locations, 2) numpy.ndarray + xy : (..., 2) numpy.ndarray receiver locations Returns ------- - numpy.ndarray + (n_t, ..., 3) numpy.ndarray Magnetic flux time derivative for each xy location """ - - dxy = xy - self.location + try: + xy = check_xyz_dim(xy, dim=2) + except ValueError: + xy = check_xyz_dim(xy, dim=3)[..., :2] + dxy = xy - self.location[:2] dh_dt = magnetic_field_time_deriv_magnetic_dipole( - np.r_[self.time], dxy[:, :2], self.sigma, self.mu, self.moment + self.time, dxy, self.sigma, self.mu, self.moment ) - return dh_dt[0] + return dh_dt def magnetic_flux_time_derivative(self, xy): """Magnetic flux due to a step off magnetic dipole over a half space @@ -86,12 +94,12 @@ def magnetic_flux_time_derivative(self, xy): Parameters ---------- - xy : (n_locations, 2) numpy.ndarray + xy : (..., 2) numpy.ndarray receiver locations Returns ------- - numpy.ndarray + (n_t, ..., 3) numpy.ndarray magnetic flux for each xy location """ return self.mu * self.magnetic_field_time_derivative(xy) diff --git a/geoana/em/tdem/meson.build b/geoana/em/tdem/meson.build index 415e68f0..46df6f72 100644 --- a/geoana/em/tdem/meson.build +++ b/geoana/em/tdem/meson.build @@ -5,6 +5,7 @@ python_sources = [ 'halfspace.py', 'simple_functions.py', 'wholespace.py', + 'reference.py', ] py.install_sources( diff --git a/geoana/em/tdem/reference.py b/geoana/em/tdem/reference.py new file mode 100644 index 00000000..bdb60596 --- /dev/null +++ b/geoana/em/tdem/reference.py @@ -0,0 +1,42 @@ +""" +Useful reference functions from Ward and Hohmann. + +These do not do anything special, and are just direct translations, with +as few modifications as possible, of specific equations from Ward and Hohmann. +""" +from scipy.special import erf, ive, iv +import numpy as np + +def hz_from_vert_4_69a(m, theta, rho): + front = m / (4 * np.pi * rho**3) + tp = theta * rho + erf_tp = erf(tp) + inner_1 = 9 / (2 * tp**2) * erf_tp + inner_2 = erf_tp + inner_3 = (9 / tp + 4 * tp)/np.sqrt(np.pi) * np.exp(-tp**2) + return front * (inner_1 - inner_2 - inner_3) + +def dhz_from_vert_4_70(m, theta, rho, sigma, mu): + front = -m / (2 * np.pi * mu * sigma * rho**5) + tp = theta * rho + erf_tp = erf(tp) + inner_1 = 9 * erf_tp + inner_2 = 2 * tp / np.sqrt(np.pi) * ( + 9 + 6*tp**2 + 4 * tp**4 + ) * np.exp(-tp**2) + return front * (inner_1 - inner_2) + +def hp_from_vert_4_72(m, theta, rho): + tp = theta * rho + first = -m * theta**2 / (2 * np.pi * rho) + second = (ive(1, tp**2/2) - ive(2, tp**2/2)) + return first * second + +def dhp_from_vert_4_74(m, theta, rho, t): + tp = theta * rho + front = m * theta**2 / (2 * np.pi * rho * t) * np.exp(-tp**2 / 2) + + inner1 = (1 + tp**2)*iv(0, tp**2 / 2) + inner2 = (2 + tp**2 + 4/tp**2) * iv(1, tp**2/2) + + return front * (inner1 - inner2) \ No newline at end of file diff --git a/geoana/em/tdem/simple_functions.py b/geoana/em/tdem/simple_functions.py index d3d66cda..43b96a84 100644 --- a/geoana/em/tdem/simple_functions.py +++ b/geoana/em/tdem/simple_functions.py @@ -1,7 +1,10 @@ +from tabnanny import check + import numpy as np from scipy.special import erf, ive from scipy.constants import mu_0 from geoana.em.tdem.base import theta +from geoana.utils import append_ndim, check_xyz_dim def vertical_magnetic_field_horizontal_loop( @@ -63,11 +66,9 @@ def vertical_magnetic_field_horizontal_loop( >>> plt.show() """ - theta = np.sqrt((sigma * mu_0) / (4 * t)) - ta = theta * radius - eta = erf(ta) + ta = theta(t, sigma, mu) * radius t1 = (3 / (np.sqrt(np.pi) * ta)) * np.exp(-(ta ** 2)) - t2 = (1 - (3 / (2 * ta ** 2))) * eta + t2 = (1 - (3 / (2 * ta ** 2))) * erf(ta) hz = (t1 + t2) / (2 * radius) return turns * current * hz @@ -167,17 +168,16 @@ def vertical_magnetic_field_time_deriv_horizontal_loop( >>> plt.ylabel(r'$\\frac{\\partial h_z}{ \\partial t}$ (A/(m s)') >>> plt.show() """ - a = radius - the = theta(t, sigma, mu) - return -turns * current / (mu * sigma * a**3) * ( - 3*erf(the * a) - 2/np.sqrt(np.pi) * the * a * (3 + 2 * the**2 * a**2) * np.exp(-the**2 * a**2) + dbz_dt = vertical_magnetic_flux_time_deriv_horizontal_loop( + t, sigma=sigma, mu=mu, radius=radius, current=current, turns=turns ) + return dbz_dt / mu def vertical_magnetic_flux_time_deriv_horizontal_loop( t, sigma=1.0, mu=mu_0, radius=1.0, current=1.0, turns=1 ): - """Time-derivative of the vertical transient magnetic flux density at the center of a horizontal loop over a halfspace. + r"""Time-derivative of the vertical transient magnetic flux density at the center of a horizontal loop over a halfspace. Compute the time-derivative of the vertical component of the transient magnetic flux density at the center of a circular loop on the surface @@ -225,14 +225,13 @@ def vertical_magnetic_flux_time_deriv_horizontal_loop( >>> plt.loglog(times*1E3, -dbz_dt, '--') >>> plt.xlabel('time (ms)') - >>> plt.ylabel(r'$\\frac{\\partial b_z}{ \\partial t}$ (T/s)') + >>> plt.ylabel(r'$\frac{\partial b_z}{ \partial t}$ (T/s)') >>> plt.show() """ a = radius - the = theta(t, sigma, mu) - return -turns * current / (sigma * a**3) * ( - 3*erf(the * a) - 2/np.sqrt(np.pi) * the * a * (3 + 2 * the**2 * a**2) * np.exp(-the**2 * a**2) - ) + ta = theta(t, sigma, mu) * a + term = 3 * erf(ta) - 2 * ta / np.sqrt(np.pi) * (3 + 2 * ta**2) * np.exp(-ta**2) + return -turns * current / (sigma * a**3) * term def magnetic_field_vertical_magnetic_dipole( @@ -244,7 +243,7 @@ def magnetic_field_vertical_magnetic_dipole( ---------- t : (n_t) numpy.ndarray times (s) - xy : (n_locs, 2) numpy.ndarray + xy : (..., 2) numpy.ndarray surface field locations (m) sigma : float, optional conductivity @@ -255,7 +254,7 @@ def magnetic_field_vertical_magnetic_dipole( Returns ------- - h : (n_t, n_locs, 3) numpy.ndarray + h : (n_t, ..., 3) numpy.ndarray The magnetic field at the observation locations and times. Notes @@ -303,10 +302,17 @@ def magnetic_field_vertical_magnetic_dipole( >>> plt.legend() >>> plt.show() """ - r = np.linalg.norm(xy[:, :2], axis=-1) - x = xy[:, 0] - y = xy[:, 1] - thr = theta(t, sigma, mu=mu)[:, None] * r #will be positive... + try: + xy = check_xyz_dim(xy, 3)[..., :2] + except ValueError: + xy = check_xyz_dim(xy, 2) + + r = np.linalg.norm(xy[..., :2], axis=-1) + x = xy[..., 0] + y = xy[..., 1] + t = append_ndim(t, r.ndim) + + thr = theta(t, sigma, mu=mu) * r #will be positive... h_z = 1.0 / r**3 * ( (9 / (2 * thr**2) - 1) * erf(thr) @@ -314,14 +320,9 @@ def magnetic_field_vertical_magnetic_dipole( ) # positive here because z+ up - # iv(1, arg) - iv(2, arg) - # ive(1, arg) * np.exp(abs(arg)) - ive(2, arg) * np.exp(abs(arg)) - # (ive(1, arg) - ive(2, arg))*np.exp(abs(arg)) h_r = 2 * thr**2 / r**3 * ( ive(1, thr**2 / 2) - ive(2, thr**2 / 2) ) - # thetar is always positive so this above simplifies (more numerically stable) - angle = np.arctan2(y, x) h_x = np.cos(angle) * h_r h_y = np.sin(angle) * h_r @@ -338,7 +339,7 @@ def magnetic_field_time_deriv_magnetic_dipole( ---------- t : (n_t) numpy.ndarray times (s) - xy : (n_locs, 2) numpy.ndarray + xy : (..., 2) numpy.ndarray surface field locations (m) sigma : float, optional conductivity @@ -349,7 +350,7 @@ def magnetic_field_time_deriv_magnetic_dipole( Returns ------- - dh_dt : (n_t, n_locs, 3) numpy.ndarray + dh_dt : (n_t, ..., 3) numpy.ndarray The magnetic field at the observation locations and times. Notes @@ -400,18 +401,23 @@ def magnetic_field_time_deriv_magnetic_dipole( >>> plt.legend() >>> plt.show() """ - r = np.linalg.norm(xy[:, :2], axis=-1) - x = xy[:, 0] - y = xy[:, 1] - tr = theta(t, sigma, mu)[:, None] * r - - dhz_dt = 1 / (r**3 * t[:, None]) * ( + try: + xy = check_xyz_dim(xy, 3)[..., :2] + except ValueError: + xy = check_xyz_dim(xy, 2) + r = np.linalg.norm(xy, axis=-1) + x = xy[..., 0] + y = xy[..., 1] + t = append_ndim(t, r.ndim) + tr = theta(t, sigma, mu) * r + + dhz_dt = 1 / (r**3 * t) * ( 9 / (2 * tr**2) * erf(tr) - (4 * tr**3 + 6 * tr + 9/tr)/np.sqrt(np.pi)*np.exp(-tr**2) ) # iv(k, v) = ive(k, v) * exp(abs(arg)) - dhr_dt = - 2 * tr**2 / (r**3 * t[:, None]) * ( + dhr_dt = -2 * tr**2 / (r**3 * t) * ( (1 + tr**2) * ive(0, tr**2 / 2) - (2 + tr**2 + 4 / tr**2) * ive(1, tr**2 / 2) ) @@ -431,7 +437,7 @@ def magnetic_flux_vertical_magnetic_dipole( ---------- t : (n_t) numpy.ndarray times (s) - xy : (n_locs, 2) numpy.ndarray + xy : (..., 2) numpy.ndarray surface field locations (m) sigma : float, optional conductivity @@ -442,7 +448,7 @@ def magnetic_flux_vertical_magnetic_dipole( Returns ------- - b : (n_t, n_locs, 3) numpy.ndarray + b : (n_t, ..., 3) numpy.ndarray The magnetic flux at the observation locations and times. See Also @@ -464,7 +470,7 @@ def magnetic_flux_time_deriv_magnetic_dipole( ---------- t : (n_t) numpy.ndarray times (s) - xy : (n_locs, 2) numpy.ndarray + xy : (..., 2) numpy.ndarray surface field locations (m) sigma : float, optional conductivity @@ -475,7 +481,7 @@ def magnetic_flux_time_deriv_magnetic_dipole( Returns ------- - db_dt : (n_t, n_locs, 3) numpy.ndarray + db_dt : (n_t, ..., 3) numpy.ndarray The magnetic flux at the observation locations and times. See Also diff --git a/geoana/em/tdem/wholespace.py b/geoana/em/tdem/wholespace.py index dbd21f5a..f862b726 100644 --- a/geoana/em/tdem/wholespace.py +++ b/geoana/em/tdem/wholespace.py @@ -1,8 +1,9 @@ from scipy.special import erf import numpy as np + from geoana.em.tdem.base import BaseTDEM from geoana.spatial import repeat_scalar -from geoana.utils import check_xyz_dim +from geoana.utils import check_xyz_dim, append_ndim from geoana.em.base import BaseElectricDipole ############################################################################### @@ -42,15 +43,14 @@ def vector_potential(self, xyz): Parameters ---------- - xyz : (n, 3) numpy.ndarray + xyz : (..., 3) numpy.ndarray Gridded xyz locations Returns ------- - (n_time, n_loc, 3) numpy.array of float + (n_time, ..., 3) numpy.ndarray of float Vector potential at all times for the gridded - locations provided. Output array is squeezed when n_time and/or - n_loc = 1. + locations provided. Examples -------- @@ -90,23 +90,15 @@ def vector_potential(self, xyz): >>> ax.set_ylabel('Z') >>> ax.set_title('Vector potential at {} s'.format(time[t_ind])) """ + xyz = check_xyz_dim(xyz) - r = self.distance(xyz) - - n_loc = len(r) - n_time = len(self.time) - - theta_r = np.outer(self.theta, r) - tile_r = np.outer(np.ones(n_time), r) - - term_1 = ( - (self.current * self.length) * erf(theta_r) / (4 * np.pi * tile_r**3) - ).reshape((n_time, n_loc, 1)) - term_1 = np.tile(term_1, (1, 1, 3)) + r_vec = xyz - self.location + r = np.linalg.norm(r_vec, axis=-1, keepdims=True) + theta = append_ndim(self.theta, r.ndim) - term_2 = np.tile(np.reshape(self.orientation, (1, 1, 3)), (n_time, n_loc, 1)) + theta_r = theta * r - return (term_1 * term_2).squeeze() + return self.current * self.length / (4 * np.pi * r) * erf(theta_r) * self.orientation def electric_field(self, xyz): r"""Electric field for the transient current dipole at a set of gridded locations. @@ -132,15 +124,14 @@ def electric_field(self, xyz): Parameters ---------- - xyz : (n, 3) numpy.ndarray + xyz : (..., 3) numpy.ndarray Gridded xyz locations Returns ------- - (n_time, n_loc, 3) numpy.array of float + (n_time, ..., 3) numpy.ndarray of float Electric field at all times for the gridded - locations provided. Output array is squeezed when n_time and/or - n_loc = 1. + locations provided. Examples -------- @@ -181,65 +172,27 @@ def electric_field(self, xyz): >>> ax.set_title('Electric field at {} s'.format(time[t_ind])) """ - # dxyz = self.vector_distance(xyz) - # r = self.distance(xyz) - # r = repeat_scalar(r) - # theta_r = self.theta * r - # root_pi = np.sqrt(np.pi) - - # front = ( - # (self.current * self.length) / (4 * np.pi * self.sigma * r**3) - # ) - - # symmetric_term = ( - # ( - # - ( - # 4/root_pi * theta_r ** 3 + 6/root_pi * theta_r - # ) * np.exp(-theta_r**2) + - # 3 * erf(theta_r) - # ) * ( - # repeat_scalar(self.dot_orientation(dxyz)) * dxyz / r**2 - # ) - # ) - - # oriented_term = ( - # ( - # 4./root_pi * theta_r**3 + 2./root_pi * theta_r - # ) * np.exp(-theta_r**2) - - # erf(theta_r) - # ) * np.kron(self.orientation, np.ones((dxyz.shape[0], 1))) - - # return front * (symmetric_term + oriented_term) + xyz = check_xyz_dim(xyz) - root_pi = np.sqrt(np.pi) - dxyz = self.vector_distance(xyz) - r = self.distance(xyz) + r_vec = xyz - self.location + r = np.linalg.norm(r_vec, axis=-1, keepdims=True) + r_hat = r_vec / r + theta = append_ndim(self.theta, r.ndim) + theta_r = theta * r - n_loc = len(r) - n_time = len(self.time) + root_pi = np.sqrt(np.pi) - theta_r = np.outer(self.theta, r) - tile_r = np.outer(np.ones(n_time), r) - r = repeat_scalar(r) + front = self.current * self.length / (4 * np.pi * self.sigma * r**3) - front = ( - (self.current * self.length) / (4 * np.pi * self.sigma * tile_r**3) - ).reshape((n_time, n_loc, 1)) - front = np.tile(front, (1, 1, 3)) + common_part = theta_r / root_pi * np.exp(-theta_r**2) - term_1 = 3 * erf(theta_r) - (4/root_pi * theta_r ** 3 + 6/root_pi * theta_r) * np.exp(-theta_r**2) - term_1 = np.tile(term_1.reshape((n_time, n_loc, 1)), (1, 1, 3)) - term_2 = repeat_scalar(self.dot_orientation(dxyz)) * dxyz / r**2 - term_2 = np.tile(term_2.reshape((1, n_loc, 3)), (n_time, 1, 1)) - symmetric_term = term_1 * term_2 + sym_term = 3 * erf(theta_r) - (4 * theta_r**2 + 6) * common_part + sym_direction = r_hat.dot(self.orientation)[..., None] * r_hat - term_1 = (4./root_pi * theta_r**3 + 2./root_pi * theta_r) * np.exp(-theta_r**2) - erf(theta_r) - term_1 = np.tile(term_1.reshape((n_time, n_loc, 1)), (1, 1, 3)) - term_2 = np.kron(self.orientation, np.ones((dxyz.shape[0], 1))) - term_2 = np.tile(term_2.reshape((1, n_loc, 3)), (n_time, 1, 1)) - oriented_term = term_1 * term_2 + orient_term = erf(theta_r) - (4 * theta_r**2 + 2) * common_part + orient_dir = self.orientation - return (front * (symmetric_term + oriented_term)).squeeze() + return front * (sym_term * sym_direction - orient_term * orient_dir) def current_density(self, xyz): r"""Current density for the transient current dipole at a set of gridded locations. @@ -265,15 +218,14 @@ def current_density(self, xyz): Parameters ---------- - xyz : (n, 3) numpy.ndarray + xyz : (..., 3) numpy.ndarray Gridded xyz locations Returns ------- - (n_time, n_loc, 3) numpy.array of float + (n_time, ..., 3) numpy.ndarray of float Current density at all times for the gridded - locations provided. Output array is squeezed when n_time and/or - n_loc = 1. + locations provided. Examples -------- @@ -339,15 +291,14 @@ def magnetic_field(self, xyz): Parameters ---------- - xyz : (n, 3) numpy.ndarray + xyz : (..., 3) numpy.ndarray Gridded xyz locations Returns ------- - (n_time, n_loc, 3) numpy.array of float + (n_time, ..., 3) numpy.ndarray of float Transient magnetic field at all times for the gridded - locations provided. Output array is squeezed when n_time and/or - n_loc = 1. + locations provided. Examples -------- @@ -388,38 +339,21 @@ def magnetic_field(self, xyz): >>> ax.set_title('Magnetic field at {} s'.format(time[t_ind])) """ - # dxyz = self.vector_distance(xyz) - # r = self.distance(dxyz) - # r = repeat_scalar(r) - # thetar = self.theta * r - - # front = ( - # self.current * self.length / (4 * np.pi * r**2) * ( - # 2 / np.sqrt(np.pi) * thetar * np.exp(-thetar**2) + erf(thetar) - # ) - # ) - - # return - front * self.cross_orientation(xyz) / r - - dxyz = self.vector_distance(xyz) - r = self.distance(dxyz) + xyz = check_xyz_dim(xyz) - n_loc = len(r) - n_time = len(self.time) + r_vec = xyz - self.location + r = np.linalg.norm(r_vec, axis=-1, keepdims=True) + r_hat = r_vec / r + theta = append_ndim(self.theta, r.ndim) - theta_r = np.outer(self.theta, r) - tile_r = np.outer(np.ones(n_time), r) + theta_r = theta * r - term_1 = ( - self.current * self.length / (4 * np.pi * tile_r**2) * ( - - 2 / np.sqrt(np.pi) * theta_r * np.exp(-theta_r**2) + erf(theta_r) - ) - ).reshape((n_time, n_loc, 1)) + term1 = (self.current * self.length) / (4 * np.pi * r**2) + term2 = erf(theta_r) - 2 / np.sqrt(np.pi) * theta_r * np.exp(-theta_r**2) - r = repeat_scalar(r) - term_2 = (self.cross_orientation(xyz) / r).reshape((1, n_loc, 3)) + direction = np.cross(self.orientation, r_hat) - return - (np.tile(term_1, (1, 1, 3)) * np.tile(term_2, (n_time, 1, 1))).squeeze() + return term1 * term2 * direction def magnetic_field_time_deriv(self, xyz): r"""Time-derivative of the magnetic field for the transient current dipole at a set of gridded locations. @@ -444,15 +378,14 @@ def magnetic_field_time_deriv(self, xyz): Parameters ---------- - xyz : (n, 3) numpy.ndarray + xyz : (..., 3) numpy.ndarray Gridded xyz locations Returns ------- - (n_time, n_loc, 3) numpy.array of float + (n_time, ..., 3) numpy.ndarray of float Time-derivative of the transient magnetic field at all times for the gridded - locations provided. Output array is squeezed when n_time and/or - n_loc = 1. + locations provided. Examples -------- @@ -493,37 +426,22 @@ def magnetic_field_time_deriv(self, xyz): >>> ax.set_title('dH/dt at {} s'.format(time[t_ind])) """ - # dxyz = self.vector_distance(xyz) - # r = self.distance(dxyz) - # r = repeat_scalar(r) - - # front = ( - # self.current * self.length * self.theta**3 * r / - # (2 * np.sqrt(np.pi)**3 * self.time) - # ) - - # return - front * self.cross_orientation(xyz) / r - - dxyz = self.vector_distance(xyz) - r = self.distance(dxyz) + xyz = check_xyz_dim(xyz) - n_loc = len(r) - n_time = len(self.time) + r_vec = xyz - self.location + r = np.linalg.norm(r_vec, axis=-1, keepdims=True) + r_hat = r_vec / r + theta = append_ndim(self.theta, r.ndim) + time = append_ndim(self.time, r.ndim) - theta_r = np.outer(self.theta, r) - theta3_r = np.outer(self.theta**3, r) - tile_t = np.outer(self.time, np.ones(n_loc)) + theta_r = theta * r - term_1 = ( - self.current * self.length * theta3_r / - (2 * np.sqrt(np.pi)**3 * tile_t) * - np.exp(-theta_r**2) - ).reshape((n_time, n_loc, 1)) + term1 = (self.current * self.length) / (4 * np.pi * r**2) + term2_dt = -2 * theta_r**3 / (np.sqrt(np.pi) * time) * np.exp(-theta_r**2) - r = repeat_scalar(r) - term_2 = (self.cross_orientation(xyz) / r).reshape((1, n_loc, 3)) + direction = np.cross(self.orientation, r_hat) - return - (np.tile(term_1, (1, 1, 3)) * np.tile(term_2, (n_time, 1, 1))).squeeze() + return term1 * term2_dt * direction def magnetic_flux_density(self, xyz): r"""Magnetic flux density for the transient current dipole at a set of gridded locations. @@ -548,15 +466,14 @@ def magnetic_flux_density(self, xyz): Parameters ---------- - xyz : (n, 3) numpy.ndarray + xyz : (..., 3) numpy.ndarray Gridded xyz locations Returns ------- - (n_time, n_loc, 3) numpy.array of float + (n_time, ..., 3) numpy.ndarray of float Transient magnetic field at all times for the gridded - locations provided. Output array is squeezed when n_time and/or - n_loc = 1. + locations provided. Examples -------- @@ -623,15 +540,14 @@ def magnetic_flux_density_time_deriv(self, xyz): Parameters ---------- - xyz : (n, 3) numpy.ndarray + xyz : (..., 3) numpy.ndarray Gridded xyz locations Returns ------- - (n_time, n_loc, 3) numpy.array of float + (n_time, ..., 3) numpy.ndarray of float Time-derivative of the transient magnetic flux density at all times for the gridded - locations provided. Output array is squeezed when n_time and/or - n_loc = 1. + locations provided. Examples -------- @@ -679,13 +595,14 @@ def magnetic_flux_density_time_deriv(self, xyz): class TransientPlaneWave(BaseTDEM): """ Class for simulating the fields for a transient planewave in a wholespace. - The wave is assumed to be propogating vertically downward. + The wave is assumed to be propogating vertically downward, starting at time=0 + and z = 0 Parameters ---------- amplitude : float amplitude of primary electric field. Default is 1 - orientation : (3) array_like or {'X','Y'} + orientation : (3,) array_like or {'X','Y'} Orientation of the planewave. Can be defined using as an ``array_like`` of length 3, or by using one of {'X','Y'} to define a planewave along the x or y direction. Default is 'X'. @@ -767,8 +684,8 @@ def electric_field(self, xyz): Returns ------- - (n_t, ..., 3) numpy.array of float - Electric field at all frequencies for the gridded + (n_t, ..., 3) numpy.ndarray of float + Electric field at all times for the gridded locations provided. Examples @@ -811,17 +728,16 @@ def electric_field(self, xyz): xyz = check_xyz_dim(xyz) e0 = self.amplitude - z = xyz[..., 2] - t = self.time - for i in range(z.ndim): - t = t[..., None] + z = xyz[..., [2]] + t = append_ndim(self.time, xyz.ndim) + mu = self.mu sigma = self.sigma bunja = -e0 * (mu * sigma) ** 0.5 * z * np.exp(-(mu * sigma * z ** 2) / (4 * t)) bunmo = 2 * np.pi ** 0.5 * t ** 1.5 - return self.orientation * (bunja / bunmo)[..., None] + return self.orientation * (bunja / bunmo) def current_density(self, xyz): r"""Current density for the transient planewave at a set of gridded locations. @@ -833,7 +749,7 @@ def current_density(self, xyz): Returns ------- - (n_t, ..., 3) numpy.array of float + (n_t, ..., 3) numpy.ndarray of float Current density at all frequencies for the gridded locations provided. @@ -886,7 +802,7 @@ def magnetic_field(self, xyz): Returns ------- - (n_t, ..., 3) numpy.array of float + (n_t, ..., 3) numpy.ndarray of float Magnetic field at all frequencies for the gridded locations provided. @@ -939,7 +855,7 @@ def magnetic_flux_density(self, xyz): Returns ------- - (n_t, ..., 3) numpy.array of float + (n_t, ..., 3) numpy.ndarray of float Magnetic flux density at all frequencies for the gridded locations provided. @@ -984,12 +900,11 @@ def magnetic_flux_density(self, xyz): e0 = self.amplitude - z = xyz[..., 2] - t = self.time - for i in range(z.ndim): - t = t[..., None] - mu = self.mu + z = xyz[..., [2]] + t = append_ndim(self.time, xyz.ndim) + sigma = self.sigma + mu = self.mu # Curl E = -dB/dt b_amp = - e0 * np.sqrt(sigma * mu / (np.pi * t)) * np.exp((-mu * sigma * z ** 2)/(4 * t)) @@ -998,4 +913,4 @@ def magnetic_flux_density(self, xyz): # take cross product with the propagation direction b_dir = np.cross(self.orientation, [0, 0, -1]) - return b_dir * b_amp[..., None] + return b_dir * b_amp diff --git a/geoana/gravity.py b/geoana/gravity.py index 71dd2d6e..64fb1eda 100644 --- a/geoana/gravity.py +++ b/geoana/gravity.py @@ -205,8 +205,9 @@ def gravitational_field(self, xyz): """ xyz = check_xyz_dim(xyz) r_vec = xyz - self.location - r = np.linalg.norm(r_vec, axis=-1) - g_vec = -G * self.mass * r_vec / r[..., None] ** 3 + r = np.linalg.norm(r_vec, axis=-1, keepdims=True) + r_hat = r_vec / r + g_vec = -G * self.mass * r_hat / r ** 2 return g_vec def gravitational_gradient(self, xyz): @@ -266,9 +267,11 @@ def gravitational_gradient(self, xyz): """ xyz = check_xyz_dim(xyz) r_vec = xyz - self.location - r = np.linalg.norm(r_vec, axis=-1) - g_tens = -G * self.mass * (np.eye(3) / r[..., None, None] ** 3 - - 3 * r_vec[..., None] * r_vec[..., None, :] / r[..., None, None] ** 5) + r = np.linalg.norm(r_vec, axis=-1, keepdims=True) + r_hat = r_vec / r + g_tens = -G * self.mass * ( + np.eye(3) - 3 * r_hat[..., None] * r_hat[..., None, :] + ) / r[..., None]**3 return g_tens diff --git a/geoana/kernels/_extensions/rTE.pyx b/geoana/kernels/_extensions/rTE.pyx index 25928d1d..b7fb64e8 100644 --- a/geoana/kernels/_extensions/rTE.pyx +++ b/geoana/kernels/_extensions/rTE.pyx @@ -1,5 +1,6 @@ # distutils: language=c++ # cython: language_level=3 +# cython: embedsignature=True import numpy as np cimport numpy as np cimport cython diff --git a/geoana/kernels/potential_field_prism.py b/geoana/kernels/potential_field_prism.py index 34173446..6d7682fa 100644 --- a/geoana/kernels/potential_field_prism.py +++ b/geoana/kernels/potential_field_prism.py @@ -1,4 +1,5 @@ import numpy as np +from scipy.special import xlogy def _prism_f(x, y, z): diff --git a/geoana/spatial.py b/geoana/spatial.py index 127be522..5c6d067f 100644 --- a/geoana/spatial.py +++ b/geoana/spatial.py @@ -34,16 +34,16 @@ def cylindrical_to_cartesian(grid, vec=None): Parameters ---------- - grid : (n, 3) array_like + grid : (..., 3) array_like Location points defined in cylindrical coordinates :math:`(r, \\phi, z)`. - vec : (n, 3) array_like, optional + vec : (..., 3) array_like, optional Vectors defined in cylindrical coordinates :math:`(v_r, v_\\phi, v_z)` at the gridded locations. Will also except a flattend array in column major order with the same number of elements. Returns ------- - (n, 3) numpy.ndarray + (..., 3) numpy.ndarray If `vec` is ``None``, this returns the transformed `grid` array, otherwise this is the transformed `vec` array. @@ -86,27 +86,25 @@ def cylindrical_to_cartesian(grid, vec=None): [ 7.07106781e-01, -7.07106781e-01, 3.00000000e+00], [ 1.00000000e+00, -2.44929360e-16, 4.00000000e+00]]) """ - grid = np.atleast_2d(grid) + grid = np.asarray(grid) if vec is None: - return np.hstack([ - mkvc(grid[:, 0]*np.cos(grid[:, 1]), 2), - mkvc(grid[:, 0]*np.sin(grid[:, 1]), 2), - mkvc(grid[:, 2], 2) - ]) - - if len(vec.shape) == 1 or vec.shape[1] == 1: - vec = vec.reshape(grid.shape, order='F') - - x = vec[:, 0] * np.cos(grid[:, 1]) - vec[:, 1] * np.sin(grid[:, 1]) - y = vec[:, 0] * np.sin(grid[:, 1]) + vec[:, 1] * np.cos(grid[:, 1]) - - newvec = [x, y] - if grid.shape[1] == 3: - z = vec[:, 2] - newvec += [z] - - return np.vstack(newvec).T + out = [ + grid[..., 0]*np.cos(grid[..., 1]), + grid[..., 0]*np.sin(grid[..., 1]), + ] + if grid.shape[-1] == 3: + out.append(grid[..., 2]) + else: + vec = np.asarray(vec) + + out = [ + vec[..., 0] * np.cos(grid[..., 1]) - vec[..., 1] * np.sin(grid[..., 1]), + vec[..., 0] * np.sin(grid[..., 1]) + vec[..., 1] * np.cos(grid[..., 1]) + ] + if grid.shape[-1] == 3: + out.append(vec[..., 2]) + return np.stack(out, axis=-1) def cartesian_to_cylindrical(grid, vec=None): @@ -118,16 +116,16 @@ def cartesian_to_cylindrical(grid, vec=None): Parameters ---------- - grid : (n, 3) array_like + grid : (..., 3) array_like Gridded locations defined in Cartesian coordinates :math:`(x, y z)`. - vec : (n, 3) array_like, optional + vec : (..., 3) array_like, optional Vectors defined in Cartesian coordinates :math:`(v_x, v_y, v_z)` at the gridded locations. Also accepts a flattened array with the same total elements in column major order. Returns ------- - (n, 3) numpy.ndarray + (..., 3) numpy.ndarray If `vec` is ``None``, this returns the transformed `grid` array, otherwise this is the transformed `vec` array. @@ -172,26 +170,28 @@ def cartesian_to_cylindrical(grid, vec=None): [ 1.00000000e+00, -7.85398163e-01, 3.00000000e+00], [ 1.00000000e+00, -2.44929360e-16, 4.00000000e+00]]) """ - - grid = np.atleast_2d(grid) + grid = np.asarray(grid) if vec is None: - return np.hstack([ - mkvc(np.sqrt(grid[:, 0]**2 + grid[:, 1]**2), 2), - mkvc(np.arctan2(grid[:, 1], grid[:, 0]), 2), - mkvc(grid[:, 2], 2) - ]) - - if len(vec.shape) == 1 or vec.shape[1] == 1: - vec = vec.reshape(grid.shape, order='F') - - theta = np.arctan2(grid[:, 1], grid[:, 0]) + out = [ + np.sqrt(grid[..., 0]**2 + grid[..., 1]**2), + np.arctan2(grid[..., 1], grid[..., 0]), + ] + if grid.shape[-1] == 3: + out.append( + grid[..., 2] + ) + else: + vec = np.asarray(vec) - return np.hstack([ - mkvc(np.cos(theta)*vec[:, 0] + np.sin(theta)*vec[:, 1], 2), - mkvc(-np.sin(theta)*vec[:, 0] + np.cos(theta)*vec[:, 1], 2), - mkvc(vec[:, 2], 2) - ]) + theta = np.arctan2(grid[..., 1], grid[..., 0]) + out = [ + np.cos(theta)*vec[..., 0] + np.sin(theta)*vec[..., 1], + -np.sin(theta) * vec[..., 0] + np.cos(theta) * vec[..., 1] + ] + if grid.shape[-1] == 3: + out.append(vec[..., 2]) + return np.stack(out, axis=-1) def spherical_to_cartesian(grid, vec=None): @@ -203,16 +203,16 @@ def spherical_to_cartesian(grid, vec=None): Parameters ---------- - grid : (n, 3) array_like + grid : (..., 3) array_like Gridded locations defined in spherical coordinates :math:`(r, \\phi, \\theta)`. - vec : (n, 3) array_like, optional + vec : (..., 3) array_like, optional Vectors defined in spherical coordinates :math:`(v_r, v_\\phi, v_\\theta)` at the gridded locations. Will also except a flattend array in column major order with the same number of elements. Returns ------- - (n, 3) numpy.ndarray + (..., 3) numpy.ndarray If `vec` is ``None``, this returns the transformed `grid` array, otherwise this is the transformed `vec` array. @@ -255,36 +255,35 @@ def spherical_to_cartesian(grid, vec=None): [ 2.70598050e-01, -2.70598050e-01, -9.23879533e-01], [ 1.22464680e-16, -2.99951957e-32, -1.00000000e+00]]) """ - grid = np.atleast_2d(grid) + grid = np.asarray(grid) if vec is None: - return np.hstack([ - mkvc(grid[:, 0] * np.sin(grid[:, 2]) * np.cos(grid[:, 1]), 2), - mkvc(grid[:, 0] * np.sin(grid[:, 2]) * np.sin(grid[:, 1]), 2), - mkvc(grid[:, 0] * np.cos(grid[:, 2]), 2) - ]) + z = grid[..., 0] * np.cos(grid[..., 2]) + rho = grid[..., 0] * np.sin(grid[..., 2]) + x = rho * np.cos(grid[..., 1]) + y = rho * np.sin(grid[..., 1]) + return np.stack([x, y, z], axis=-1) - if len(vec.shape) == 1 or vec.shape[1] == 1: - vec = vec.reshape(grid.shape, order='F') + vec = np.asarray(vec) + + phi = grid[..., 1] + theta = grid[..., 2] x = ( - vec[:, 0] * np.sin(grid[:, 2]) * np.cos(grid[:, 1]) + - vec[:, 2] * np.cos(grid[:, 2]) * np.cos(grid[:, 1]) - - vec[:, 1] * np.sin(grid[:, 1]) + vec[..., 0] * np.sin(theta) * np.cos(phi) - + vec[..., 1] * np.sin(phi) + + vec[..., 2] * np.cos(theta) * np.cos(phi) ) + y = ( - vec[:, 0] * np.sin(grid[:, 2]) * np.sin(grid[:, 1]) + - vec[:, 2] * np.cos(grid[:, 2]) * np.sin(grid[:, 1]) - - vec[:, 1] * np.cos(grid[:, 1]) - ) - z = ( - vec[:, 0] * np.cos(grid[:, 2]) - - vec[:, 2] * np.sin(grid[:, 2]) + vec[..., 0] * np.sin(theta) * np.sin(phi) + + vec[..., 1] * np.cos(theta) + + vec[..., 2] * np.cos(theta) * np.sin(phi) ) - newvec = [x, y, z] + z = vec[..., 0] * np.cos(theta) - vec[..., 2] * np.sin(theta) - return np.vstack(newvec).T + return np.stack([x, y, z], axis=-1) def cartesian_to_spherical(grid, vec=None): @@ -296,16 +295,16 @@ def cartesian_to_spherical(grid, vec=None): Parameters ---------- - grid : (n, 3) array_like + grid : (..., 3) array_like Gridded locations defined in Cartesian coordinates :math:`(x, y, z)`. - vec : (n, 3) array_like, optional + vec : (..., 3) array_like, optional Vectors defined in Cartesian coordinates :math:`(v_x, v_y, v_z)` at the gridded locations. Will also except a flattend array in column major order with the same number of elements. Returns ------- - (n, 3) numpy.ndarray + (..., 3) numpy.ndarray If `vec` is ``None``, this returns the transformed `grid` array, otherwise this is the transformed `vec` array. @@ -352,41 +351,35 @@ def cartesian_to_spherical(grid, vec=None): [ 1.00000000e+00, -2.44929360e-16, 3.14159265e+00]]) """ - grid = np.atleast_2d(grid) + # phi is azimuth + # theta is polar - if vec is None: - return np.hstack([ - mkvc(np.sqrt(grid[:, 0]**2 + grid[:, 1]**2 + grid[:, 2]**2), 2), - mkvc(np.arctan2(grid[:, 1], grid[:, 0]), 2), - mkvc( - np.arctan2(np.sqrt(grid[:, 0]**2 + grid[:, 1]**2), grid[:, 2]), - 2 - ), - ]) + grid = np.asarray(grid) + rho = np.linalg.norm(grid[..., :2], axis=-1) + phi = np.arctan2(grid[..., 1], grid[..., 0]) + theta = np.arctan2(rho, grid[..., 2]) - if len(vec.shape) == 1 or vec.shape[1] == 1: - vec = vec.reshape(grid.shape, order='F') + if vec is None: + r = np.linalg.norm(grid, axis=-1) + return np.stack([r, phi, theta], axis=-1) - theta = np.arctan2(grid[:, 1], grid[:, 0]) - phi = np.arctan2(np.sqrt(grid[:, 0]**2 + grid[:, 1]**2), grid[:, 2]) + vec = np.asarray(vec) r = ( - vec[:, 0] * np.sin(phi) * np.cos(theta) + - vec[:, 1] * np.sin(phi) * np.sin(theta) + - vec[:, 2] * np.cos(phi) + vec[..., 0] * np.sin(theta) * np.cos(phi) + + vec[..., 1] * np.sin(theta) * np.sin(phi) + + vec[..., 2] * np.cos(theta) ) - theta = - vec[:, 0] * np.sin(theta) + vec[:, 1] * np.cos(theta) + phi = - vec[..., 0] * np.sin(phi) + vec[..., 1] * np.cos(phi) - phi = ( - vec[:, 0] * np.cos(phi) * np.cos(theta) + - vec[:, 1] * np.cos(phi) * np.sin(theta) - - vec[:, 2] * np.sin(phi) + theta = ( + vec[..., 0] * np.cos(theta) * np.cos(phi) + + vec[..., 1] * np.cos(theta) * np.sin(phi) - + vec[..., 2] * np.sin(theta) ) - newvec = [r, theta, phi] - - return np.vstack(newvec).T + return np.stack([r, phi, theta], axis=-1) def vector_magnitude(v): @@ -394,18 +387,15 @@ def vector_magnitude(v): Parameters ---------- - (n, dim) numpy.ndarray + (..., dim) numpy.ndarray A set of input vectors (2D or 3D) Returns ------- - (n) numpy.ndarray + (...) numpy.ndarray Magnitudes of the vectors """ - - v = np.atleast_2d(v) - - return np.sqrt((v**2).sum(axis=1)) + return np.linalg.norm(v, axis=-1) def vector_distance(xyz, origin=np.r_[0., 0., 0.]): @@ -497,7 +487,7 @@ def vector_dot(xyz, vector): ---------- xyz : (n, 3) numpy.ndarray A set of 3D vectors - vector : (3) numpy.array_like + vector : (3) numpy.ndarray_like A single 3D vector Returns diff --git a/geoana/utils.py b/geoana/utils.py index 101714d6..574531cf 100644 --- a/geoana/utils.py +++ b/geoana/utils.py @@ -17,7 +17,6 @@ """ import numpy as np - def mkvc(x, n_dims=1): """Creates a vector with specified dimensionality. @@ -182,6 +181,13 @@ def check_xyz_dim(xyz, dim=3, dtype=float): return xyz +def append_ndim(arr, ndim): + arr = np.asarray(arr) + for _ in range(ndim): + arr = arr[..., None] + return arr + + def requires(modules): """Decorate a function with soft dependencies. @@ -221,4 +227,4 @@ def passer(*args, **kwargs): return passer - return decorated_function \ No newline at end of file + return decorated_function diff --git a/pyproject.toml b/pyproject.toml index fd878079..3c005348 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -105,3 +105,14 @@ setup-args = [ '-Dwith_extensions=true', '--vsenv' ] + +[tool.pytest.ini_options] +filterwarnings = [ + 'ignore:divide by zero encountered in divide:RuntimeWarning', + 'ignore:invalid value encountered in divide:RuntimeWarning', + 'ignore:invalid value encountered in multiply:RuntimeWarning', + 'ignore:invalid value encountered in sqrt:RuntimeWarning', + 'ignore::numpy.exceptions.ComplexWarning' +] + + diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 00000000..0ebd0091 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,327 @@ +import pytest +import numpy as np + +from sympy.vector import CoordSys3D, Del +import sympy +from scipy.constants import gravitational_constant as gamma + +c = 299792458 # m/s +mu_0 = (4 * sympy.pi * sympy.sympify(10)**-7) +epsilon_0 = 1/(mu_0 * c**2) + +def vector_lambdify(vec_func, coord_sys, *args): + + vec_func = vec_func.to_matrix(coord_sys) + lambs = [sympy.lambdify((*args, *coord_sys.base_scalars()), f) for f in vec_func] + + def out(*inner_args): + out_shape = np.broadcast(*inner_args).shape + expansion = np.ones(out_shape) + outs = [lamb(*inner_args) * expansion for lamb in lambs] + return np.stack(outs, axis=-1) + return out + +def tensor_lambdify(row_vec_funcs, coord_sys, *args): + lambs = [] + for vec_func in row_vec_funcs: + vec_func = vec_func.to_matrix(coord_sys) + lambs.append([sympy.lambdify((*args, *coord_sys.base_scalars()), f) for f in vec_func]) + + def out(*inner_args): + out_shape = np.broadcast(*inner_args).shape + expansion = np.ones(out_shape) + outs = [] + for row in lambs: + out_row = [lamb(*inner_args) * expansion for lamb in row] + outs.append(np.stack(out_row, axis=-1)) + return np.stack(outs, axis=-1) + return out + +@pytest.fixture(scope='session') +def em_dipole_params(): + param_dict = { + "current" : sympy.sympify(2), + "length" : sympy.sympify(1)/2, + "mu" : 2 * mu_0, + "sigma" : sympy.sympify(10)**-3, + "epsilon" : 4 * epsilon_0, + "moment" : sympy.sympify(1)/2 + } + return param_dict + + +@pytest.fixture(scope='session') +def grav_point_params(): + param_dict = { + "mass" : sympy.sympify(10)**5, + "radius" : sympy.sympify(10), + } + return param_dict + + +@pytest.fixture(scope='session') +def sympy_tdem_ex_dipole(em_dipole_params): + R = CoordSys3D('R') + delop = Del() + + I = em_dipole_params['current'] + ds = em_dipole_params['length'] + mu = em_dipole_params['mu'] + sigma = em_dipole_params['sigma'] + t = sympy.symbols('t') + + r = sympy.sqrt(R.x**2 + R.y**2 + R.z**2) + theta = sympy.sqrt(mu * sigma / (4 * t)) + + A = I * ds / (4 * sympy.pi * r) * sympy.erf(theta * r) * R.i + H = delop.cross(A).doit().simplify() + B = mu * H + dH_dt = sympy.diff(H, t).doit().simplify() + dB_dt = mu * dH_dt + E = (delop.cross(H).doit() / sigma).simplify() + J = sigma * E + + lamb_funcs = { + 'vector_potential' : vector_lambdify(A, R, t), + 'magnetic_field' : vector_lambdify(H, R, t), + 'magnetic_flux_density' : vector_lambdify(B, R, t), + 'electric_field' : vector_lambdify(E, R, t), + 'current_density' : vector_lambdify(J, R, t), + 'magnetic_field_time_deriv' : vector_lambdify(dH_dt, R, t), + 'magnetic_flux_density_time_deriv' : vector_lambdify(dB_dt, R, t), + } + return lamb_funcs + + +@pytest.fixture(scope='session') +def sympy_fdem_ex_dipole(em_dipole_params): + R = CoordSys3D('R') + delop = Del() + + I = em_dipole_params['current'] + ds = em_dipole_params['length'] + mu = em_dipole_params['mu'] + sigma = em_dipole_params['sigma'] + eps = em_dipole_params['epsilon'] + f = sympy.symbols(r'f') + w = 2 * sympy.pi * f + + r = sympy.sqrt(R.x**2 + R.y**2 + R.z**2) + k = sympy.sqrt(w**2 * mu * eps - sympy.I * w * mu * sigma) + + A = I * ds / (4 * sympy.pi * r) * sympy.exp(-sympy.I * k * r) * R.i + H = delop.cross(A).doit().simplify() + B = mu * H + E = (delop.cross(H).doit() / sigma).simplify() + J = sigma * E + + lamb_funcs = { + 'vector_potential' : vector_lambdify(A, R, f), + 'magnetic_field' : vector_lambdify(H, R, f), + 'magnetic_flux_density' : vector_lambdify(B, R, f), + 'electric_field' : vector_lambdify(E, R, f), + 'current_density' : vector_lambdify(J, R, f), + } + return lamb_funcs + + +@pytest.fixture(scope='session') +def sympy_fdem_hx_dipole(em_dipole_params): + + R = CoordSys3D('R') + delop = Del() + + moment = em_dipole_params['moment'] + mu = em_dipole_params['mu'] + sigma = em_dipole_params['sigma'] + eps = em_dipole_params['epsilon'] + f = sympy.symbols(r'f') + w = 2 * sympy.pi * f + + r = sympy.sqrt(R.x**2 + R.y**2 + R.z**2) + k = sympy.sqrt(w**2 * mu * eps - sympy.I * w * mu * sigma) + + F = sympy.I * w * mu * moment / (4 * sympy.pi * r) * sympy.exp(-sympy.I * k * r) * R.i + E = -delop.cross(F).doit().simplify() + J = sigma * E + B = delop.cross(E).doit().simplify()/(-sympy.I * w) + H = B/mu + + lamb_funcs = { + 'vector_potential' : vector_lambdify(F, R, f), + 'magnetic_field' : vector_lambdify(H, R, f), + 'magnetic_flux_density' : vector_lambdify(B, R, f), + 'electric_field' : vector_lambdify(E, R, f), + 'current_density' : vector_lambdify(J, R, f), + } + return lamb_funcs + + +@pytest.fixture(scope='session') +def sympy_static_hx_dipole(em_dipole_params): + R = CoordSys3D('R') + delop = Del() + + moment = em_dipole_params['moment'] * R.i + mu = em_dipole_params['mu'] + + r_vec = R.x * R.i + R.y * R.j + R.z * R.k + r = sympy.sqrt(R.x**2 + R.y**2 + R.z**2) + + A = (mu / (4 * sympy.pi * r**3) * moment.cross(r_vec)).doit() + B = delop.cross(A).doit().simplify() + H = B / mu + + lamb_funcs = { + 'vector_potential' : vector_lambdify(A, R), + 'magnetic_field' : vector_lambdify(H, R), + 'magnetic_flux_density' : vector_lambdify(B, R), + } + return lamb_funcs + +@pytest.fixture(scope='session') +def sympy_static_h_pole(em_dipole_params): + R = CoordSys3D('R') + delop = Del() + + moment = em_dipole_params['moment'] + mu = em_dipole_params['mu'] + + r = sympy.sqrt(R.x ** 2 + R.y ** 2 + R.z ** 2) + mag_potential = moment / (4 * sympy.pi * r) + + H = -delop.gradient(mag_potential).doit() + B = mu * H + + lamb_funcs = { + 'potential' : sympy.lambdify((R.x, R.y, R.z), mag_potential), + 'magnetic_field' : vector_lambdify(H, R), + 'magnetic_flux_density' : vector_lambdify(B, R), + } + return lamb_funcs + + +@pytest.fixture(scope='session') +def sympy_grav_point(grav_point_params): + R = CoordSys3D('R') + delop = Del() + mass = grav_point_params['mass'] + + r = sympy.sqrt(R.x ** 2 + R.y ** 2 + R.z ** 2) + grav_potential = gamma * mass / r + + grav_vector = delop.gradient(grav_potential).doit() + + grav_grads = [delop.gradient(grav_vector.components[comp]).doit() for comp in R.base_vectors()] + + lamb_funcs = { + 'gravitational_potential': sympy.lambdify((R.x, R.y, R.z), grav_potential), + 'gravitational_field': vector_lambdify(grav_vector, R), + 'gravitational_gradient': tensor_lambdify(grav_grads, R), + } + + return lamb_funcs + +@pytest.fixture(scope='session') +def sympy_grav_sphere(grav_point_params): + R = CoordSys3D('R') + delop = Del() + mass = grav_point_params['mass'] + radius = grav_point_params['radius'] + volume = 4 * sympy.pi * radius**3 / 3 + density = mass/volume + + r = sympy.sqrt(R.x ** 2 + R.y ** 2 + R.z ** 2) + grav_potential = sympy.Piecewise( + (gamma * 2 * sympy.pi * density * (radius**2 - r**2/3), r < radius), + (gamma * mass / r, True), + ) + + grav_vector = delop.gradient(grav_potential).doit() + + grav_grads = [delop.gradient(grav_vector.components[comp]).doit() for comp in R.base_vectors()] + + lamb_funcs = { + 'gravitational_potential': sympy.lambdify((R.x, R.y, R.z), grav_potential), + 'gravitational_field': vector_lambdify(grav_vector, R), + 'gravitational_gradient': tensor_lambdify(grav_grads, R), + } + + return lamb_funcs + +@pytest.fixture(scope='session') +def sympy_potential_prism(): + x, y, z = sympy.symbols('x y z') + r = sympy.sqrt(x ** 2 + y ** 2 + z ** 2) + f = ( + - x * y * sympy.log(z + r) + - y * z * sympy.log(x + r) + - z * x * sympy.log(y + r) + + x ** 2 * sympy.atan(y * z / (x * r)) / 2 + + y ** 2 * sympy.atan(x * z / (y * r)) / 2 + + z ** 2 * sympy.atan(x * y / (z * r)) / 2 + ) + + fz = -sympy.diff(f, z) + + fzx = sympy.diff(f, z, x) + fzy = sympy.diff(f, z, y) + fzz = sympy.diff(f, z, z) + + fzzz = -sympy.diff(f, z, z, z) + fxxy = -sympy.diff(f, x, x, y) + fxxz = -sympy.diff(f, x, x, z) + fxyz = -sympy.diff(f, x, y, z) + + lamb_funcs = { + 'prism_f': sympy.lambdify((x, y, z), f), + 'prism_fz': sympy.lambdify((x, y, z), fz), + 'prism_fzx': sympy.lambdify((x, y, z), fzx), + 'prism_fzy': sympy.lambdify((x, y, z), fzy), + 'prism_fzz': sympy.lambdify((x, y, z), fzz), + 'prism_fzzz': sympy.lambdify((x, y, z), fzzz), + 'prism_fxxy': sympy.lambdify((x, y, z), fxxy), + 'prism_fxxz': sympy.lambdify((x, y, z), fxxz), + 'prism_fxyz': sympy.lambdify((x, y, z), fxyz), + } + + return lamb_funcs + +@pytest.fixture(scope='session') +def sympy_linex_segment(em_dipole_params): + R = CoordSys3D('R') + delop = Del() + + l = em_dipole_params['length'] + mu = em_dipole_params['mu'] + I = em_dipole_params['current'] + sigma = em_dipole_params['sigma'] + + x, y, z = R.base_scalars() + r0 = sympy.sqrt(x ** 2 + y ** 2 + z ** 2) + x1 = x - l + r1 = sympy.sqrt(x1 ** 2 + y ** 2 + z ** 2) + + A = -mu * I / (4 * sympy.pi) * ( + sympy.log(x1 + r1) - sympy.log(x + r0) + ) * R.i + + B = delop.cross(A).doit() + H = B / mu + + phi = I / (4 * sympy.pi * sigma ) * (1/r1 - 1/r0) + E = -delop.gradient(phi).doit() + J = sigma * E + + lamb_funcs = { + 'scalar_potential': sympy.lambdify(R.base_scalars(), phi), + 'vector_potential' : vector_lambdify(A, R), + 'magnetic_field' : vector_lambdify(H, R), + 'magnetic_flux_density' : vector_lambdify(B, R), + 'electric_field' : vector_lambdify(E, R), + 'current_density' : vector_lambdify(J, R), + + } + return lamb_funcs + + diff --git a/tests/test_em_fdem_base.py b/tests/em/fdem/test_em_fdem_base.py similarity index 67% rename from tests/test_em_fdem_base.py rename to tests/em/fdem/test_em_fdem_base.py index e1ef7e42..bc3225ef 100644 --- a/tests/test_em_fdem_base.py +++ b/tests/em/fdem/test_em_fdem_base.py @@ -26,17 +26,19 @@ def test_sigma_hat(): def test_base_fdem(): - edws = fdem.ElectricDipoleWholeSpace(1) - with pytest.raises(TypeError): - edws.frequency = "string" - with pytest.raises(ValueError): - edws.frequency = -1 - with pytest.raises(TypeError): - edws.frequency = np.array([[1, 2], [3, 4]]) + # gets frequency and quasistatic properties. + with pytest.raises(TypeError, match="frequencies are not a valid type"): + fdem.BaseFDEM("string") + with pytest.raises(ValueError, match="All frequencies must be greater than 0"): + fdem.BaseFDEM(-1) + with pytest.raises(TypeError, match="frequencies must have at most 1 dimension."): + fdem.BaseFDEM(np.array([[1, 2], [3, 4]])) fd = fdem.BaseFDEM(frequency=np.logspace(1, 4, 3), sigma=1, quasistatic=True) fds = fd.sigma_hat np.testing.assert_equal(1, fds) + assert fd.quasistatic is True + diff --git a/tests/em/fdem/test_em_fdem_electric_dipole.py b/tests/em/fdem/test_em_fdem_electric_dipole.py new file mode 100644 index 00000000..39feaff5 --- /dev/null +++ b/tests/em/fdem/test_em_fdem_electric_dipole.py @@ -0,0 +1,96 @@ +import pytest +import numpy as np +import numpy.testing as npt + +from geoana.em.fdem import ElectricDipoleWholeSpace + +METHODS = [ + 'vector_potential', + 'magnetic_field', + 'electric_field', + 'current_density', + 'magnetic_flux_density', +] + +@pytest.fixture() +def e_dipole(em_dipole_params): + nf = 5 + + current = float(em_dipole_params['current']) + length = float(em_dipole_params['length']) + mu = float(em_dipole_params['mu']) + conductivity = float(em_dipole_params['sigma']) + eps = float(em_dipole_params['epsilon']) + + frequencies = np.logspace(2, 5, nf) + dip = ElectricDipoleWholeSpace( + frequency=frequencies, + sigma=conductivity, + mu=mu, + current=current, + orientation='x', + length=length, + epsilon=eps, + quasistatic=False + ) + return dip + + +@pytest.fixture() +def xyz(): + nx, ny, nz = (2, 3, 4) + x = np.linspace(-5, 5, nx) + y = np.linspace(-5, 5, ny) + z = np.linspace(-5, 5, nz) + xyz = np.meshgrid(x, y, z) + return xyz + +@pytest.mark.parametrize('method', METHODS) +@pytest.mark.parametrize('nf', [0, 1, 5]) +def test_broadcasting(e_dipole, xyz, method, nf): + if nf == 0: + e_dipole.frequency = e_dipole.frequency[0] + else: + e_dipole.frequency = e_dipole.frequency[:nf] + dipole_func = getattr(e_dipole, method) + freq_shape = e_dipole.frequency.shape + + out = dipole_func(xyz) + assert out.shape == (*freq_shape, *xyz[0].shape, 3) + + xyz = np.stack(xyz, axis=-1) + out = dipole_func(xyz) + assert out.shape == (*freq_shape, *xyz.shape[:-1], 3) + + xyz = xyz.reshape((-1, 3)) + out = dipole_func(xyz) + assert out.shape == (*freq_shape, xyz.shape[0], 3) + + xyz = xyz.reshape((-1, 3)) + out = dipole_func(xyz) + assert out.shape == (*freq_shape, xyz.shape[0], 3) + + out = dipole_func(xyz[0]) + assert out.shape == (*freq_shape, 3) + + +@pytest.mark.parametrize('method', METHODS) +@pytest.mark.parametrize('orient', ['x', 'y', 'z']) +def test_correct(method, orient, e_dipole, xyz, sympy_fdem_ex_dipole): + x, y, z = xyz + e_dipole.orientation = orient + out = getattr(e_dipole, method)(xyz) + + sympy_func = sympy_fdem_ex_dipole[method] + + # cycle the input and output if not x + if orient == 'x': + verify = sympy_func(e_dipole.frequency[:, None, None, None], x, y, z) + if orient == 'y': + verify = sympy_func(e_dipole.frequency[:, None, None, None], y, z, x) + verify = verify[..., [2, 0, 1]] + elif orient == 'z': + verify = sympy_func(e_dipole.frequency[:, None, None, None], z, x, y) + verify = verify[..., [1, 2, 0]] + + npt.assert_allclose(verify, out) diff --git a/tests/test_em_fdem_layered.py b/tests/em/fdem/test_em_fdem_layered.py similarity index 100% rename from tests/test_em_fdem_layered.py rename to tests/em/fdem/test_em_fdem_layered.py diff --git a/tests/em/fdem/test_em_fdem_magnetic_dipole.py b/tests/em/fdem/test_em_fdem_magnetic_dipole.py new file mode 100644 index 00000000..6b8c18fd --- /dev/null +++ b/tests/em/fdem/test_em_fdem_magnetic_dipole.py @@ -0,0 +1,96 @@ +import pytest +import numpy as np +import numpy.testing as npt + +from geoana.em.fdem import MagneticDipoleWholeSpace + +METHODS = [ + 'vector_potential', + 'magnetic_field', + 'electric_field', + 'current_density', + 'magnetic_flux_density', +] + + +@pytest.fixture() +def h_dipole(em_dipole_params): + nf = 5 + + moment = float(em_dipole_params['moment']) + mu = float(em_dipole_params['mu']) + conductivity = float(em_dipole_params['sigma']) + eps = float(em_dipole_params['epsilon']) + + frequencies = np.logspace(2, 5, nf) + dip = MagneticDipoleWholeSpace( + frequency=frequencies, + sigma=conductivity, + mu=mu, + moment=moment, + orientation='x', + epsilon=eps, + quasistatic=False + ) + return dip + + +@pytest.fixture() +def xyz(): + nx, ny, nz = (2, 3, 4) + x = np.linspace(-5, 5, nx) + y = np.linspace(-5, 5, ny) + z = np.linspace(-5, 5, nz) + xyz = np.meshgrid(x, y, z) + return xyz + + +@pytest.mark.parametrize('method', METHODS) +@pytest.mark.parametrize('nf', [0, 1, 5]) +def test_broadcasting(h_dipole, xyz, method, nf): + if nf == 0: + h_dipole.frequency = h_dipole.frequency[0] + else: + h_dipole.frequency = h_dipole.frequency[:nf] + dipole_func = getattr(h_dipole, method) + freq_shape = h_dipole.frequency.shape + + out = dipole_func(xyz) + assert out.shape == (*freq_shape, *xyz[0].shape, 3) + + xyz = np.stack(xyz, axis=-1) + out = dipole_func(xyz) + assert out.shape == (*freq_shape, *xyz.shape[:-1], 3) + + xyz = xyz.reshape((-1, 3)) + out = dipole_func(xyz) + assert out.shape == (*freq_shape, xyz.shape[0], 3) + + xyz = xyz.reshape((-1, 3)) + out = dipole_func(xyz) + assert out.shape == (*freq_shape, xyz.shape[0], 3) + + out = dipole_func(xyz[0]) + assert out.shape == (*freq_shape, 3) + + +@pytest.mark.parametrize('method', METHODS) +@pytest.mark.parametrize('orient', ['x', 'y', 'z']) +def test_correct(method, orient, h_dipole, xyz, sympy_fdem_hx_dipole): + x, y, z = xyz + h_dipole.orientation = orient + out = getattr(h_dipole, method)(xyz) + + sympy_func = sympy_fdem_hx_dipole[method] + + # cycle the input and output if not x + if orient == 'x': + verify = sympy_func(h_dipole.frequency[:, None, None, None], x, y, z) + if orient == 'y': + verify = sympy_func(h_dipole.frequency[:, None, None, None], y, z, x) + verify = verify[..., [2, 0, 1]] + elif orient == 'z': + verify = sympy_func(h_dipole.frequency[:, None, None, None], z, x, y) + verify = verify[..., [1, 2, 0]] + + npt.assert_allclose(verify, out) diff --git a/tests/test_em_fdem_planewave.py b/tests/em/fdem/test_em_fdem_planewave.py similarity index 96% rename from tests/test_em_fdem_planewave.py rename to tests/em/fdem/test_em_fdem_planewave.py index 03ad0fb3..c5669fc4 100644 --- a/tests/test_em_fdem_planewave.py +++ b/tests/em/fdem/test_em_fdem_planewave.py @@ -140,7 +140,7 @@ def test_magnetic_field(): hy = -1 / Z * np.exp(ikz) hz = np.zeros_like(z) - h_vec = hpw.magnetic_field(xyz)[0] + h_vec = hpw.magnetic_field(xyz) np.testing.assert_equal(hx, h_vec[:, 0]) np.testing.assert_allclose(hy, h_vec[:, 1], rtol=1E-15) @@ -153,7 +153,7 @@ def test_magnetic_field(): hy = np.zeros_like(z) hz = np.zeros_like(z) - h_vec = hpw.magnetic_field(xyz)[0] + h_vec = hpw.magnetic_field(xyz) np.testing.assert_allclose(hx, h_vec[:, 0], rtol=1E-15) np.testing.assert_equal(hy, h_vec[:, 1]) @@ -171,7 +171,7 @@ def test_magnetic_flux_density(): y = np.linspace(-30., 30., 50) z = np.linspace(-40., 40., 50) xyz = np.stack(np.meshgrid(x, y, z), axis=-1).reshape(-1, 3) - z = xyz[:, 2] + z = xyz[..., 2] kz = z * k ikz = 1j * kz @@ -181,7 +181,7 @@ def test_magnetic_flux_density(): by = -1 / Z * np.exp(ikz) bz = np.zeros_like(z) - b_vec = hpw.magnetic_flux_density(xyz)[0] + b_vec = hpw.magnetic_flux_density(xyz) np.testing.assert_equal(bx, b_vec[:, 0]) np.testing.assert_allclose(by, b_vec[:, 1], rtol=1E-15) @@ -194,7 +194,7 @@ def test_magnetic_flux_density(): by = np.zeros_like(z) bz = np.zeros_like(z) - b_vec = hpw.magnetic_flux_density(xyz)[0] + b_vec = hpw.magnetic_flux_density(xyz) np.testing.assert_allclose(bx, b_vec[:, 0], rtol=1E-15) np.testing.assert_equal(by, b_vec[:, 1]) diff --git a/tests/em/static/test_current_loop.py b/tests/em/static/test_current_loop.py new file mode 100644 index 00000000..b9bb08e0 --- /dev/null +++ b/tests/em/static/test_current_loop.py @@ -0,0 +1,94 @@ +import pytest +import numpy as np +import numpy.testing as npt +from scipy.constants import mu_0 + +from geoana.em.static import CircularLoopWholeSpace, LineCurrentWholeSpace +from geoana.spatial import cylindrical_2_cartesian + +METHODS = [ + 'vector_potential', 'magnetic_field', 'magnetic_flux_density' +] + +@pytest.fixture() +def circle_loop(): + + radius = 5 + current = 1 + + sigma = 1E-3 + mu = 4 * mu_0 + loop = CircularLoopWholeSpace(radius=radius, current=current, sigma=sigma, mu=mu) + + return loop + +@pytest.fixture() +def xyz(): + nx, ny, nz = (11, 13, 9) + x = np.linspace(-100, 100, nx) + y = np.linspace(-90, 90, ny) + z = np.linspace(-80, 80, nz) + xyz = np.meshgrid(x, y, z) + return xyz + + +@pytest.mark.parametrize('method', METHODS) +def test_broadcasting(circle_loop, xyz, method): + + func = getattr(circle_loop, method) + + out = func(xyz) + assert out.shape == (*xyz[0].shape, 3) + + xyz = np.stack(xyz, axis=-1) + out = func(xyz) + assert out.shape == (*xyz.shape[:-1], 3) + + xyz = xyz.reshape((-1, 3)) + out = func(xyz) + assert out.shape == (xyz.shape[0], 3) + + xyz = xyz.reshape((-1, 3)) + out = func(xyz) + assert out.shape == (xyz.shape[0], 3) + + out = func(xyz[0]) + assert out.shape == (3,) + + +@pytest.mark.parametrize('orient', ['x', 'y', 'z']) +@pytest.mark.parametrize('method', METHODS) +def test_circular_loop(circle_loop, orient, method, xyz): + circle_loop.orientation = orient + + # line segment approx. + n_segments = 256 + rs = np.full(n_segments + 1, circle_loop.radius) + thetas = np.linspace(0, 2 * np.pi, n_segments + 1) + + # defined as if orient = 'z' + nodes = cylindrical_2_cartesian(np.c_[rs, thetas, np.zeros(n_segments + 1)]) + # Ensure it wraps. + nodes[-1] = nodes[0] + + # cycle node for other orientations + if orient == 'x': + nodes = np.c_[nodes[:, 2], nodes[:, 0], nodes[:, 1]] + elif orient == 'y': + nodes = np.c_[nodes[:, 1], nodes[:, 2], nodes[:, 0]] + + loop_approx = LineCurrentWholeSpace( + nodes, + current=circle_loop.current, + sigma=circle_loop.sigma, + mu=circle_loop.mu + ) + + test = getattr(circle_loop, method)(xyz) + other = getattr(loop_approx, method)(xyz) + + atol = 1E-16 + if method == 'magnetic_field': + atol /= mu_0 + + npt.assert_allclose(test, other, rtol=1E-3, atol=atol) \ No newline at end of file diff --git a/tests/test_em_static.py b/tests/em/static/test_em_static.py similarity index 58% rename from tests/test_em_static.py rename to tests/em/static/test_em_static.py index 96a0533f..30b6700a 100644 --- a/tests/test_em_static.py +++ b/tests/em/static/test_em_static.py @@ -2,14 +2,9 @@ import unittest import numpy as np from scipy.constants import mu_0, epsilon_0 -try: - import discretize -except ImportError: - discretize = None -from scipy.special import ellipk, ellipe +import numpy.testing as npt -from geoana.em import static, fdem -from geoana import spatial +from geoana.em import static TOL = 0.1 @@ -17,48 +12,30 @@ class TestEM_Static(unittest.TestCase): def setUp(self): - self.mdws = static.MagneticDipoleWholeSpace() self.mpws = static.MagneticPoleWholeSpace() self.clws = static.CircularLoopWholeSpace() self.lcfs = static.LineCurrentFreeSpace(nodes=np.c_[1., 1., 1.]) def test_defaults(self): - self.assertTrue(self.mdws.sigma == 1.0) self.assertTrue(self.mpws.sigma == 1.0) self.assertTrue(self.clws.sigma == 1.0) - self.assertTrue(self.mdws.mu == mu_0) self.assertTrue(self.mpws.mu == mu_0) self.assertTrue(self.clws.mu == mu_0) - self.assertTrue(self.mdws.epsilon == epsilon_0) self.assertTrue(self.mpws.epsilon == epsilon_0) self.assertTrue(self.clws.epsilon == epsilon_0) - self.assertTrue(np.all(self.mdws.orientation == np.r_[1., 0., 0.])) self.assertTrue(np.all(self.mpws.orientation == np.r_[1., 0., 0.])) self.assertTrue(np.all(self.clws.orientation == np.r_[1., 0., 0.])) self.assertTrue(np.all(self.lcfs.nodes == np.c_[1., 1., 1.])) - self.assertTrue(self.mdws.moment == 1.0) self.assertTrue(self.mpws.moment == 1.0) self.assertTrue(self.clws.current == 1.0) self.assertTrue(self.clws.radius == np.sqrt(1/np.pi)) def test_errors(self): - with pytest.raises(TypeError): - self.mdws.mu = "box" - with pytest.raises(ValueError): - self.mdws.mu = -2 - with pytest.raises(TypeError): - self.mdws.moment = "box" - with pytest.raises(ValueError): - self.mdws.orientation = [0, 1, 2, 3, 4] - with pytest.raises(ValueError): - self.mdws.orientation = [[0, 0], [0, 1]] - with pytest.raises(TypeError): - self.mdws.orientation = ["string"] with pytest.raises(ValueError): self.clws.location = [0, 1, 2, 3] with pytest.raises(ValueError): @@ -86,362 +63,6 @@ def test_errors(self): with pytest.raises(TypeError): self.lcfs.current = "box" - def test_vector_potential(self): - n = 50 - x = y = z = np.linspace(-25, 25, n) - xyz = np.stack(np.meshgrid(x, y, z), axis=-1).reshape(-1, 3) - - for orientation in ["x", "y", "z"]: - self.clws.orientation = orientation - self.mdws.orientation = orientation - - inds = ( - (np.absolute(xyz[:, 0]) > 5) & - (np.absolute(xyz[:, 1]) > 5) & - (np.absolute(xyz[:, 2]) > 5) - ) - - a_clws = self.clws.vector_potential(xyz)[inds] - a_mdws = self.mdws.vector_potential(xyz)[inds] - - self.assertTrue(isinstance(a_clws, np.ndarray)) - self.assertTrue(isinstance(a_mdws, np.ndarray)) - - self.assertTrue( - np.linalg.norm(a_clws - a_mdws) < - 0.5 * TOL * (np.linalg.norm(a_clws) + np.linalg.norm(a_mdws)) - ) - - if discretize is not None: - def test_magnetic_field_tensor(self): - print("\n === Testing Tensor Mesh === \n") - n = 30 - h = 2. - mesh = discretize.TensorMesh( - [np.full(n, h), np.full(n, h), np.full(n, h)], - origin='CCC' - ) - - for radius in [0.5, 1, 1.5]: - self.clws.radius = radius - self.clws.current = 1./(np.pi * radius**2) - - fdem_dipole = fdem.MagneticDipoleWholeSpace(frequency=0) - - for location in [ - np.r_[0, 0, 0], np.r_[10, 0, 0], np.r_[0, -10, 0], - np.r_[0, 0, 10], np.r_[10, 10, 10] - ]: - self.clws.location = location - self.mdws.location = location - fdem_dipole.location = location - - for orientation in ["x", "y", "z"]: - self.clws.orientation = orientation - self.mdws.orientation = orientation - fdem_dipole.orientation = orientation - - a_clws = np.hstack([ - self.clws.vector_potential(mesh.gridEx)[:, 0], - self.clws.vector_potential(mesh.gridEy)[:, 1], - self.clws.vector_potential(mesh.gridEz)[:, 2] - ]) - - a_mdws = np.hstack([ - self.mdws.vector_potential(mesh.gridEx)[:, 0], - self.mdws.vector_potential(mesh.gridEy)[:, 1], - self.mdws.vector_potential(mesh.gridEz)[:, 2] - ]) - - b_clws = mesh.edge_curl * a_clws - b_mdws = mesh.edge_curl * a_mdws - - b_clws_true = np.hstack([ - self.clws.magnetic_flux_density(mesh.gridFx)[:, 0], - self.clws.magnetic_flux_density(mesh.gridFy)[:, 1], - self.clws.magnetic_flux_density(mesh.gridFz)[:, 2] - ]) - b_mdws_true = np.hstack([ - self.mdws.magnetic_flux_density(mesh.gridFx)[:, 0], - self.mdws.magnetic_flux_density(mesh.gridFy)[:, 1], - self.mdws.magnetic_flux_density(mesh.gridFz)[:, 2] - ]) - b_fdem = np.hstack([ - fdem_dipole.magnetic_flux_density(mesh.gridFx)[:, 0], - fdem_dipole.magnetic_flux_density(mesh.gridFy)[:, 1], - fdem_dipole.magnetic_flux_density(mesh.gridFz)[:, 2] - ]) - - self.assertTrue(isinstance(b_fdem, np.ndarray)) - - inds = (np.hstack([ - (np.absolute(mesh.gridFx[:, 0]) > h*2 + location[0]) & - (np.absolute(mesh.gridFx[:, 1]) > h*2 + location[1]) & - (np.absolute(mesh.gridFx[:, 2]) > h*2 + location[2]), - (np.absolute(mesh.gridFy[:, 0]) > h*2 + location[0]) & - (np.absolute(mesh.gridFy[:, 1]) > h*2 + location[1]) & - (np.absolute(mesh.gridFy[:, 2]) > h*2 + location[2]), - (np.absolute(mesh.gridFz[:, 0]) > h*2 + location[0]) & - (np.absolute(mesh.gridFz[:, 1]) > h*2 + location[1]) & - (np.absolute(mesh.gridFz[:, 2]) > h*2 + location[2]) - ])) - - loop_passed_1 = ( - np.linalg.norm(b_fdem[inds] - b_clws[inds]) < - 0.5 * TOL * ( - np.linalg.norm(b_fdem[inds]) + - np.linalg.norm(b_clws[inds]) - ) - ) - - loop_passed_2 = ( - np.linalg.norm(b_fdem[inds] - b_clws_true[inds]) < - 0.5 * TOL * ( - np.linalg.norm(b_fdem[inds]) + - np.linalg.norm(b_clws_true[inds]) - ) - ) - - dipole_passed = ( - np.linalg.norm(b_fdem[inds] - b_mdws[inds]) < - 0.5 * TOL * ( - np.linalg.norm(b_fdem[inds]) + - np.linalg.norm(b_mdws[inds]) - ) - ) - - print( - "Testing r = {}, loc = {}, orientation = {}".format( - radius, location, orientation - ) - ) - print( - " fdem: {:1.4e}, loop: {:1.4e}, dipole: {:1.4e}" - " Passed? loop: {}, dipole: {} \n".format( - np.linalg.norm(b_fdem[inds]), - np.linalg.norm(b_clws[inds]), - np.linalg.norm(b_mdws[inds]), - loop_passed_1, - loop_passed_2, - dipole_passed - ) - ) - self.assertTrue(loop_passed_1) - self.assertTrue(loop_passed_2) - self.assertTrue(dipole_passed) - - def test_dipole_magnetic_flux_density(self): - mdws = static.MagneticDipoleWholeSpace() - x = np.linspace(-20., 20., 50) - y = np.linspace(-30., 30., 50) - z = np.linspace(-40., 40., 50) - xyz = np.stack(np.meshgrid(x, y, z), axis=-1).reshape(-1, 3) - n_obs = xyz.shape[0] - - xyz_ = spatial.cylindrical_2_cartesian(xyz) - r = mdws.vector_distance(xyz_) - dxyz = spatial.repeat_scalar(mdws.distance(xyz_)) - m_vec = mdws.moment * np.atleast_2d(mdws.orientation).repeat(n_obs, axis=0) - m_dot_r = (m_vec * r).sum(axis=1) - m_dot_r = np.atleast_2d(m_dot_r).T.repeat(3, axis=1) - - b_test = (mu_0 / (4 * np.pi)) * ((3 * r * m_dot_r / (dxyz ** 5)) - m_vec / (dxyz ** 3)) - b_test = spatial.cartesian_2_cylindrical(xyz_, b_test) - - b = mdws.magnetic_flux_density(xyz, coordinates='cylindrical') - np.testing.assert_equal(b_test, b) - - def test_pole_magnetic_flux_density(self): - mpws = static.MagneticPoleWholeSpace() - x = np.linspace(-20., 20., 50) - y = np.linspace(-30., 30., 50) - z = np.linspace(-40., 40., 50) - xyz = np.stack(np.meshgrid(x, y, z), axis=-1).reshape(-1, 3) - - xyz_ = spatial.cylindrical_2_cartesian(xyz) - r = mpws.vector_distance(xyz_) - dxyz = spatial.repeat_scalar(mpws.distance(xyz_)) - - b_test = mpws.moment * mu_0 / (4 * np.pi * (dxyz ** 3)) * r - b_test = spatial.cartesian_2_cylindrical(xyz_, b_test) - - b = mpws.magnetic_flux_density(xyz, coordinates='cylindrical') - np.testing.assert_equal(b_test, b) - - def test_pole_magnetic_field(self): - mpws = static.MagneticPoleWholeSpace() - x = np.linspace(-20., 20., 50) - y = np.linspace(-30., 30., 50) - z = np.linspace(-40., 40., 50) - xyz = np.stack(np.meshgrid(x, y, z), axis=-1).reshape(-1, 3) - - xyz_ = spatial.cylindrical_2_cartesian(xyz) - r = mpws.vector_distance(xyz_) - dxyz = spatial.repeat_scalar(mpws.distance(xyz_)) - - h_test = mpws.moment * mu_0 / (4 * np.pi * (dxyz ** 3)) * r - h_test = spatial.cartesian_2_cylindrical(xyz_, h_test) - h_test = h_test / mu_0 - - h = mpws.magnetic_field(xyz, coordinates='cylindrical') - np.testing.assert_equal(h_test, h) - - def test_loop_magnetic_flux_density(self): - clws = static.CircularLoopWholeSpace() - x = np.linspace(-20., 20., 50) - y = np.linspace(-30., 30., 50) - z = np.linspace(-40., 40., 50) - xyz = np.stack(np.meshgrid(x, y, z), axis=-1).reshape(-1, 3) - - xyz_ = np.atleast_2d(xyz) - xyz_ = spatial.cylindrical_2_cartesian(xyz_) - xyz_ = spatial.rotate_points_from_normals(xyz_, np.array(clws.orientation), np.r_[0, 0, 1], - x0=np.array(clws.location)) - dxyz = clws.vector_distance(xyz_) - rho = np.linalg.norm(dxyz[:, :2], axis=-1) - b = np.zeros((len(rho), 3)) - ind_axial = rho == 0 - b[ind_axial, -1] = mu_0 * clws.current * clws.radius ** 2 / (2 * (clws.radius ** 2 + - dxyz[ind_axial, 2] ** 2) ** (1.5)) - alpha = rho[~ind_axial] / clws.radius - beta = dxyz[~ind_axial, 2] / clws.radius - gamma = dxyz[~ind_axial, 2] / rho[~ind_axial] - q = ((1 + alpha) ** 2 + beta ** 2) - k2 = 4 * alpha / q - b[~ind_axial, -1] = mu_0 * clws.current / (2 * clws.radius * np.pi * np.sqrt(q)) * \ - (ellipe(k2) * (1 - alpha ** 2 - beta ** 2) / (q - 4 * alpha) + ellipk(k2)) - b_rad = mu_0 * clws.current * gamma / (2 * clws.radius * np.pi * np.sqrt(q)) * \ - (ellipe(k2) * (1 + alpha ** 2 + beta ** 2) / (q - 4 * alpha) - ellipk(k2)) - b[~ind_axial, 0] = b_rad * (dxyz[~ind_axial, 0] / rho[~ind_axial]) - b[~ind_axial, 1] = b_rad * (dxyz[~ind_axial, 1] / rho[~ind_axial]) - b = spatial.rotate_points_from_normals(b, np.r_[0, 0, 1], np.array(clws.orientation),) - b = spatial.cartesian_2_cylindrical(xyz_, b) - - b_test = clws.magnetic_flux_density(xyz, coordinates='cylindrical') - np.testing.assert_equal(b_test, b) - - if discretize is not None: - def test_magnetic_field_3Dcyl(self): - print("\n === Testing 3D Cyl Mesh === \n") - n = 50 - ny = 10 - h = 2. - mesh = discretize.CylindricalMesh( - [h*np.ones(n), np.ones(ny) * 2 * np.pi / ny, h*np.ones(n)], x0="00C" - ) - - for radius in [0.5, 1, 1.5]: - self.clws.radius = radius - self.clws.current = 1./(np.pi * radius**2) # So that all have moment of 1.0 - - fdem_dipole = fdem.MagneticDipoleWholeSpace(frequency=0, quasistatic=True) - - for location in [ - np.r_[0, 0, 0], np.r_[0, 4, 0], np.r_[4, 4., 0], - np.r_[4, 4, 4], np.r_[4, 0, -4] - ]: - self.clws.location = location - self.mdws.location = location - fdem_dipole.location = location - - for orientation in ["z"]: - self.clws.orientation = orientation - self.mdws.orientation = orientation - fdem_dipole.orientation = orientation - - a_clws = np.hstack([ - self.clws.vector_potential( - mesh.gridEx, coordinates="cylindrical" - )[:, 0], - self.clws.vector_potential( - mesh.gridEy, coordinates="cylindrical" - )[:, 1], - self.clws.vector_potential( - mesh.gridEz, coordinates="cylindrical" - )[:, 2] - ]) - - a_mdws = np.hstack([ - self.mdws.vector_potential( - mesh.gridEx, coordinates="cylindrical" - )[:, 0], - self.mdws.vector_potential( - mesh.gridEy, coordinates="cylindrical" - )[:, 1], - self.mdws.vector_potential( - mesh.gridEz, coordinates="cylindrical" - )[:, 2] - ]) - - b_clws = mesh.edge_curl * a_clws - b_mdws = mesh.edge_curl * a_mdws - - b_fdem = np.hstack([ - spatial.cartesian_2_cylindrical( - spatial.cylindrical_2_cartesian(mesh.gridFx), - fdem_dipole.magnetic_flux_density( - spatial.cylindrical_2_cartesian(mesh.gridFx) - ) - )[:, 0], - spatial.cartesian_2_cylindrical( - spatial.cylindrical_2_cartesian(mesh.gridFy), - fdem_dipole.magnetic_flux_density( - spatial.cylindrical_2_cartesian(mesh.gridFy) - ) - )[:, 1], - spatial.cartesian_2_cylindrical( - spatial.cylindrical_2_cartesian(mesh.gridFz), - fdem_dipole.magnetic_flux_density( - spatial.cylindrical_2_cartesian(mesh.gridFz) - ) - )[:, 2] - ]) - - - inds = (np.hstack([ - (np.absolute(mesh.gridFx[:, 0]) > h*4 + location[0]) & - (np.absolute(mesh.gridFx[:, 2]) > h*4 + location[2]), - (np.absolute(mesh.gridFy[:, 0]) > h*4 + location[0]) & - (np.absolute(mesh.gridFy[:, 2]) > h*4 + location[2]), - (np.absolute(mesh.gridFz[:, 0]) > h*4 + location[0]) & - (np.absolute(mesh.gridFz[:, 2]) > h*4 + location[2]) - ])) - - loop_passed = ( - np.linalg.norm(b_fdem[inds] - b_clws[inds]) < - 0.5 * TOL * ( - np.linalg.norm(b_fdem[inds]) + - np.linalg.norm(b_clws[inds]) - ) - ) - - dipole_passed = ( - np.linalg.norm(b_fdem[inds] - b_mdws[inds]) < - 0.5 * TOL * ( - np.linalg.norm(b_fdem[inds]) + - np.linalg.norm(b_mdws[inds]) - ) - ) - - print( - "Testing r = {}, loc = {}, orientation = {}".format( - radius, location, orientation - ) - ) - print( - " fdem: {:1.4e}, loop: {:1.4e}, dipole: {:1.4e}" - " Passed? loop: {}, dipole: {} \n".format( - np.linalg.norm(b_fdem[inds]), - np.linalg.norm(b_clws[inds]), - np.linalg.norm(b_mdws[inds]), - loop_passed, - dipole_passed - ) - ) - self.assertTrue(loop_passed) - self.assertTrue(dipole_passed) - - def Vt_from_ESphere( XYZ, loc, sig_s, sig_b, radius, amp ): @@ -644,14 +265,14 @@ def testV(self): vt = ess.potential(xyz, field='total') vp = ess.potential(xyz, field='primary') vs = ess.potential(xyz, field='secondary') - np.testing.assert_equal(vttest, vt) - np.testing.assert_equal(vptest, vp) - np.testing.assert_equal(vstest, vs) + npt.assert_allclose(vttest, vt) + npt.assert_allclose(vptest, vp) + npt.assert_allclose(vstest, vs) vt, vp, vs = ess.potential(xyz, field='all') - np.testing.assert_equal(vttest, vt) - np.testing.assert_equal(vptest, vp) - np.testing.assert_equal(vstest, vs) + npt.assert_allclose(vttest, vt) + npt.assert_allclose(vptest, vp) + npt.assert_allclose(vstest, vs) def testE(self): radius = 1.0 @@ -687,14 +308,14 @@ def testE(self): et = ess.electric_field(xyz, field='total') ep = ess.electric_field(xyz, field='primary') es = ess.electric_field(xyz, field='secondary') - np.testing.assert_equal(ettest, et) - np.testing.assert_equal(eptest, ep) - np.testing.assert_equal(estest, es) + npt.assert_allclose(ettest, et) + npt.assert_allclose(eptest, ep) + npt.assert_allclose(estest, es) et, ep, es =ess.electric_field(xyz, field='all') - np.testing.assert_equal(ettest, et) - np.testing.assert_equal(eptest, ep) - np.testing.assert_equal(estest, es) + npt.assert_allclose(ettest, et) + npt.assert_allclose(eptest, ep) + npt.assert_allclose(estest, es) def testJ(self): radius = 1.0 @@ -730,14 +351,14 @@ def testJ(self): jt = ess.current_density(xyz, field='total') jp = ess.current_density(xyz, field='primary') js = ess.current_density(xyz, field='secondary') - np.testing.assert_equal(jttest, jt) - np.testing.assert_equal(jptest, jp) - np.testing.assert_equal(jstest, js) + npt.assert_allclose(jttest, jt) + npt.assert_allclose(jptest, jp) + npt.assert_allclose(jstest, js) jt, jp, js = ess.current_density(xyz, field='all') - np.testing.assert_equal(jttest, jt) - np.testing.assert_equal(jptest, jp) - np.testing.assert_equal(jstest, js) + npt.assert_allclose(jttest, jt) + npt.assert_allclose(jptest, jp) + npt.assert_allclose(jstest, js) def test_rho(self): radius = 1.0 @@ -762,7 +383,7 @@ def test_rho(self): ) rho = ess.charge_density(xyz) - np.testing.assert_equal(rho_test, rho) + npt.assert_allclose(rho_test, rho) def Vt_from_Sphere( @@ -947,14 +568,14 @@ def testV(self): vt = mss.potential(xyz, field='total') vp = mss.potential(xyz, field='primary') vs = mss.potential(xyz, field='secondary') - np.testing.assert_equal(vttest, vt) - np.testing.assert_equal(vptest, vp) - np.testing.assert_equal(vstest, vs) + npt.assert_allclose(vttest, vt) + npt.assert_allclose(vptest, vp) + npt.assert_allclose(vstest, vs) vt, vp, vs = mss.potential(xyz, field='all') - np.testing.assert_equal(vttest, vt) - np.testing.assert_equal(vptest, vp) - np.testing.assert_equal(vstest, vs) + npt.assert_allclose(vttest, vt) + npt.assert_allclose(vptest, vp) + npt.assert_allclose(vstest, vs) def testH(self): radius = 1.0 @@ -990,14 +611,14 @@ def testH(self): ht = mss.magnetic_field(xyz, field='total') hp = mss.magnetic_field(xyz, field='primary') hs = mss.magnetic_field(xyz, field='secondary') - np.testing.assert_equal(httest, ht) - np.testing.assert_equal(hptest, hp) - np.testing.assert_equal(hstest, hs) + npt.assert_allclose(httest, ht) + npt.assert_allclose(hptest, hp) + npt.assert_allclose(hstest, hs) ht, hp, hs = mss.magnetic_field(xyz, field='all') - np.testing.assert_equal(httest, ht) - np.testing.assert_equal(hptest, hp) - np.testing.assert_equal(hstest, hs) + npt.assert_allclose(httest, ht) + npt.assert_allclose(hptest, hp) + npt.assert_allclose(hstest, hs) def testB(self): radius = 1.0 @@ -1033,146 +654,15 @@ def testB(self): bt = mss.magnetic_flux_density(xyz, field='total') bp = mss.magnetic_flux_density(xyz, field='primary') bs = mss.magnetic_flux_density(xyz, field='secondary') - np.testing.assert_equal(btest, bt) - np.testing.assert_equal(bptest, bp) - np.testing.assert_equal(bstest, bs) + npt.assert_allclose(btest, bt) + npt.assert_allclose(bptest, bp) + npt.assert_allclose(bstest, bs) bt, bp, bs = mss.magnetic_flux_density(xyz, field='all') - np.testing.assert_equal(btest, bt) - np.testing.assert_equal(bptest, bp) - np.testing.assert_equal(bstest, bs) - - -def V_from_PointCurrentW( - XYZ, loc, rho, cur -): - - XYZ = np.atleast_2d(XYZ) - - r_vec = XYZ - loc - r = np.linalg.norm(r_vec, axis=-1) - - v = rho * cur / (4 * np.pi * r) - return v - - -def J_from_PointCurrentW( - XYZ, loc, rho, cur -): - - j = E_from_PointCurrentW(XYZ, loc, rho, cur) / rho - return j - - -def E_from_PointCurrentW( - XYZ, loc, rho, cur -): - - XYZ = np.atleast_2d(XYZ) - - r_vec = XYZ - loc - r = np.linalg.norm(r_vec, axis=-1) - - e = rho * cur * r_vec / (4 * np.pi * r[..., None] ** 3) - return e - - -class TestPointCurrentWholeSpace: - - def test_defaults(self): - rho = 1.0 - pcws = static.PointCurrentWholeSpace(rho) - assert pcws.rho == 1.0 - assert pcws.current == 1.0 - assert np.all(pcws.location == np.r_[0., 0., 0.]) - - def test_error(self): - pcws = static.PointCurrentWholeSpace(rho=1.0, current=1.0, location=None) - - with pytest.raises(TypeError): - pcws.rho = "string" - with pytest.raises(ValueError): - pcws.rho = -1 - with pytest.raises(TypeError): - pcws.current = "string" - with pytest.raises(ValueError): - pcws.location = [0, 1, 2, 3] - with pytest.raises(ValueError): - pcws.location = [[0, 0], [0, 1]] - with pytest.raises(TypeError): - pcws.location = ["string"] - - def test_potential(self): - rho = 1.0 - current = 1.0 - location = None - pcws = static.PointCurrentWholeSpace( - current=current, - rho=rho, - location=location - ) - x = np.linspace(-20., 20., 50) - y = np.linspace(-30., 30., 50) - z = np.linspace(-40., 40., 50) - xyz = np.stack(np.meshgrid(x, y, z), axis=-1).reshape(-1, 3) - - vtest = V_from_PointCurrentW( - xyz, pcws.location, pcws.rho, pcws.current - ) - print( - "\n\nTesting Electric Potential V for Point Current\n" - ) - - v = pcws.potential(xyz) - np.testing.assert_equal(vtest, v) - - def test_current_density(self): - rho = 1.0 - current = 1.0 - location = None - pcws = static.PointCurrentWholeSpace( - current=current, - rho=rho, - location=location - ) - x = np.linspace(-20., 20., 50) - y = np.linspace(-30., 30., 50) - z = np.linspace(-40., 40., 50) - xyz = np.stack(np.meshgrid(x, y, z), axis=-1).reshape(-1, 3) - - jtest = J_from_PointCurrentW( - xyz, pcws.location, pcws.rho, pcws.current - ) - print( - "\n\nTesting Current Density J for Point Current\n" - ) - - j = pcws.current_density(xyz) - np.testing.assert_equal(jtest, j) - - def test_electric_field(self): - rho = 1.0 - current = 1.0 - location = None - pcws = static.PointCurrentWholeSpace( - current=current, - rho=rho, - location=location - ) - x = np.linspace(-20., 20., 50) - y = np.linspace(-30., 30., 50) - z = np.linspace(-40., 40., 50) - xyz = np.stack(np.meshgrid(x, y, z), axis=-1).reshape(-1, 3) - - etest = E_from_PointCurrentW( - xyz, pcws.location, pcws.rho, pcws.current - ) - print( - "\n\nTesting Electric Field E for Point Current\n" - ) + npt.assert_allclose(btest, bt) + npt.assert_allclose(bptest, bp) + npt.assert_allclose(bstest, bs) - e = pcws.electric_field(xyz) - np.testing.assert_equal(etest, e) def V_from_PointCurrentH( @@ -1279,7 +769,7 @@ def test_potential(self): ) v = pchs.potential(xyz) - np.testing.assert_equal(vtest, v) + npt.assert_allclose(vtest, v) x = np.linspace(-20., 20., 50) y = np.linspace(-30., 30., 50) @@ -1311,7 +801,7 @@ def test_electric_field(self): ) e = pchs.electric_field(xyz) - np.testing.assert_equal(etest, e) + npt.assert_allclose(etest, e) x = np.linspace(-20., 20., 50) y = np.linspace(-30., 30., 50) @@ -1343,7 +833,7 @@ def test_current_density(self): ) j = pchs.current_density(xyz) - np.testing.assert_equal(jtest, j) + npt.assert_allclose(jtest, j) x = np.linspace(-20., 20., 50) y = np.linspace(-30., 30., 50) @@ -1520,8 +1010,8 @@ def test_potential(self): v1 = dhs.potential(xyz1) v2 = dhs.potential(xyz1, xyz2) - np.testing.assert_equal(vtest1, v1) - np.testing.assert_equal(vtest2, v2) + npt.assert_allclose(vtest1, v1) + npt.assert_allclose(vtest2, v2) x = np.linspace(-20., 20., 50) y = np.linspace(-30., 30., 50) @@ -1560,8 +1050,8 @@ def test_electric_field(self): e1 = dhs.electric_field(xyz1) e2 = dhs.electric_field(xyz1, xyz2) - np.testing.assert_equal(etest1, e1) - np.testing.assert_equal(etest2, e2) + npt.assert_allclose(etest1, e1) + npt.assert_allclose(etest2, e2) x = np.linspace(-20., 20., 50) y = np.linspace(-30., 30., 50) @@ -1600,8 +1090,8 @@ def test_current_density(self): j1 = dhs.current_density(xyz1) j2 = dhs.current_density(xyz1, xyz2) - np.testing.assert_equal(jtest1, j1) - np.testing.assert_equal(jtest2, j2) + npt.assert_allclose(jtest1, j1) + npt.assert_allclose(jtest2, j2) x = np.linspace(-20., 20., 50) y = np.linspace(-30., 30., 50) diff --git a/tests/em/static/test_em_static_wholespace.py b/tests/em/static/test_em_static_wholespace.py new file mode 100644 index 00000000..3c9d8f48 --- /dev/null +++ b/tests/em/static/test_em_static_wholespace.py @@ -0,0 +1,269 @@ +import numpy as np +import pytest +import numpy.testing as npt +from scipy.constants import point, gravitational_constant, mu_0 + +from geoana.em.static import MagneticDipoleWholeSpace, MagneticPoleWholeSpace +from geoana.em.static.wholespace import LineCurrentWholeSpace, PointCurrentWholeSpace + +METHODS = [ + 'vector_potential', + 'magnetic_field', + 'magnetic_flux_density', +] + +ELEC_METHODS = [ + "scalar_potential", + "electric_field", + "current_density", +] + +@pytest.fixture() +def h_dipole(em_dipole_params): + moment = float(em_dipole_params['moment']) + mu = float(em_dipole_params['mu']) + + dip = MagneticDipoleWholeSpace( + mu=mu, + moment=moment, + orientation='x', + ) + return dip + +@pytest.fixture() +def h_pole(em_dipole_params): + moment = float(em_dipole_params['moment']) + mu = float(em_dipole_params['mu']) + + dip = MagneticPoleWholeSpace( + mu=mu, + moment=moment, + ) + return dip + +@pytest.fixture() +def line_current(em_dipole_params): + length = float(em_dipole_params['length']) + current = float(em_dipole_params['current']) + mu = float(em_dipole_params['mu']) + sigma = float(em_dipole_params['sigma']) + + nodes = np.asarray([ + [0, 0, 0,], + [length, 0, 0], + ]) + + line = LineCurrentWholeSpace(nodes, current, mu=mu, sigma=sigma) + return line + + +@pytest.fixture() +def point_current(em_dipole_params): + current = float(em_dipole_params['current']) + sigma = float(em_dipole_params['sigma']) + + location = [0, 0, 0] + + point = PointCurrentWholeSpace(1/sigma, current, location) + return point + + +@pytest.fixture() +def xyz(): + nx, ny, nz = (2, 3, 4) + x = np.linspace(-5, 5, nx) + y = np.linspace(-5, 5, ny) + z = np.linspace(-5, 5, nz) + xyz = np.meshgrid(x, y, z) + return xyz + + +@pytest.mark.parametrize('method', METHODS) +def test_h_dip_broadcasting(h_dipole, xyz, method): + dipole_func = getattr(h_dipole, method) + + out = dipole_func(xyz) + assert out.shape == (*xyz[0].shape, 3) + + xyz = np.stack(xyz, axis=-1) + out = dipole_func(xyz) + assert out.shape == (*xyz.shape[:-1], 3) + + xyz = xyz.reshape((-1, 3)) + out = dipole_func(xyz) + assert out.shape == (xyz.shape[0], 3) + + xyz = xyz.reshape((-1, 3)) + out = dipole_func(xyz) + assert out.shape == (xyz.shape[0], 3) + + out = dipole_func(xyz[0]) + assert out.shape == (3,) + + +@pytest.mark.parametrize('method', METHODS) +@pytest.mark.parametrize('orient', ['x', 'y', 'z']) +def test_h_dip_correct(method, orient, h_dipole, xyz, sympy_static_hx_dipole): + x, y, z = xyz + h_dipole.orientation = orient + out = getattr(h_dipole, method)(xyz) + + sympy_func = sympy_static_hx_dipole[method] + + # cycle the input and output if not x + if orient == 'x': + verify = sympy_func(x, y, z) + if orient == 'y': + verify = sympy_func(y, z, x) + verify = verify[..., [2, 0, 1]] + elif orient == 'z': + verify = sympy_func(z, x, y) + verify = verify[..., [1, 2, 0]] + npt.assert_allclose(verify, out, atol=1E-20) + +@pytest.mark.parametrize('method', METHODS[1:]) +def test_h_pole_broadcasting(h_pole, xyz, method): + func = getattr(h_pole, method) + + out = func(xyz) + assert out.shape == (*xyz[0].shape, 3) + + xyz = np.stack(xyz, axis=-1) + out = func(xyz) + assert out.shape == (*xyz.shape[:-1], 3) + + xyz = xyz.reshape((-1, 3)) + out = func(xyz) + assert out.shape == (xyz.shape[0], 3) + + xyz = xyz.reshape((-1, 3)) + out = func(xyz) + assert out.shape == (xyz.shape[0], 3) + + out = func(xyz[0]) + assert out.shape == (3,) + +@pytest.mark.parametrize('method', METHODS[1:]) +def test_h_pole_correct(method, h_pole, xyz, sympy_static_h_pole): + x, y, z = xyz + out = getattr(h_pole, method)(xyz) + + sympy_func = sympy_static_h_pole[method] + + verify = sympy_func(x, y, z) + npt.assert_allclose(out, verify) + +@pytest.mark.parametrize('method', METHODS + ELEC_METHODS) +def test_line_broadcasting(line_current, xyz, method): + func = getattr(line_current, method) + + if method == 'scalar_potential': + out_shape = tuple() + else: + out_shape = (3, ) + + out = func(xyz) + assert out.shape == (*xyz[0].shape, *out_shape) + + xyz = np.stack(xyz, axis=-1) + out = func(xyz) + assert out.shape == (*xyz.shape[:-1], *out_shape) + + xyz = xyz.reshape((-1, 3)) + out = func(xyz) + assert out.shape == (xyz.shape[0], *out_shape) + + xyz = xyz.reshape((-1, 3)) + out = func(xyz) + assert out.shape == (xyz.shape[0], *out_shape) + + out = func(xyz[0]) + assert out.shape == out_shape + +@pytest.mark.parametrize('method', METHODS + ELEC_METHODS) +@pytest.mark.parametrize('orient', ['x', 'y', 'z']) +def test_line_correct(method, orient, line_current, xyz, sympy_linex_segment): + x, y, z = xyz + # they all start out as oriented in the x-direction + if orient == 'y': + l = line_current.nodes[1, 0] + line_current.nodes[1] = [0, l, 0] + elif orient == 'z': + l = line_current.nodes[1, 0] + line_current.nodes[1] = [0, 0, l] + + out = getattr(line_current, method)(xyz) + + sympy_func = sympy_linex_segment[method] + + # cycle the input and output if not x + if orient == 'x': + verify = sympy_func(x, y, z) + if orient == 'y': + verify = sympy_func(y, z, x) + if method != 'scalar_potential': + verify = verify[..., [2, 0, 1]] + elif orient == 'z': + verify = sympy_func(z, x, y) + if method != 'scalar_potential': + verify = verify[..., [1, 2, 0]] + + atol = 1E-17 + if method != 'magnetic_field': + atol *= mu_0 # to account for mu + + npt.assert_allclose(out, verify, atol=atol) + + +@pytest.mark.parametrize('method', ['potential', 'electric_field', 'current_density']) +def test_point_current_broadcast(point_current, xyz, method): + func = getattr(point_current, method) + + if method == 'potential': + out_shape = tuple() + else: + out_shape = (3,) + + out = func(xyz) + assert out.shape == (*xyz[0].shape, *out_shape) + + xyz = np.stack(xyz, axis=-1) + out = func(xyz) + assert out.shape == (*xyz.shape[:-1], *out_shape) + + xyz = xyz.reshape((-1, 3)) + out = func(xyz) + assert out.shape == (xyz.shape[0], *out_shape) + + xyz = xyz.reshape((-1, 3)) + out = func(xyz) + assert out.shape == (xyz.shape[0], *out_shape) + + out = func(xyz[0]) + assert out.shape == out_shape + + +@pytest.mark.parametrize('method', ['potential', 'electric_field', 'current_density']) +def test_point_current_correct(method, xyz, point_current, sympy_grav_point): + + # scale from gravity to electric... + scale = -point_current.current/(4 * np.pi * point_current.sigma * gravitational_constant * 10**5) + if method == 'potential': + scale *= -1 + elif method == 'current_density': + scale *= point_current.sigma + + x, y, z = xyz + + grav_method = { + 'potential' : 'gravitational_potential', + 'electric_field' : 'gravitational_field', + 'current_density' : 'gravitational_field', + }[method] + + out = getattr(point_current, method)(xyz) + + verify = scale * sympy_grav_point[grav_method](x, y, z) + + npt.assert_allclose(out, verify) + diff --git a/tests/em/static/test_freespace.py b/tests/em/static/test_freespace.py new file mode 100644 index 00000000..c102f4be --- /dev/null +++ b/tests/em/static/test_freespace.py @@ -0,0 +1,54 @@ +import numpy as np +import numpy.testing as npt + +import pytest +from scipy.special import roots_legendre + +from geoana.em.static import MagneticPrism, MagneticDipoleWholeSpace +from geoana.utils import append_ndim + + +class TestMagneticAccuracy(): + x, y, z = np.mgrid[-100:100:20j, -100:100:20j, -100:100:20j] + xyz = np.stack((x, y, z), axis=-1).reshape((-1, 3)) + dx = 0.1 + prism = MagneticPrism(dx * np.r_[-1, -1, -1], dx * np.r_[1, 1, 1], magnetization=[-1, 2, -0.5]) + m_mag = np.linalg.norm(prism.magnetization) + m_unit = prism.magnetization/m_mag + dipole = MagneticDipoleWholeSpace( + location=[0, 0, 0], moment=m_mag, orientation=m_unit + ) + + quad_points, quad_weights = roots_legendre(5) + quad_points = (prism.max_location - prism.min_location)[:, None] * (quad_points + 1) / 2 + prism.min_location[:, + None] + quad_points = np.stack(np.meshgrid(*quad_points, indexing='ij'), axis=-1) + quad_wx, quad_wy, quad_wz = quad_weights * (prism.max_location - prism.min_location)[:, None] / 2 + quad_wx = quad_wx[:, None, None] + quad_wy = quad_wy[None, :, None] + quad_wz = quad_wz[None, None, :] + + quad_xyzs = xyz - quad_points[..., None, :] + + @pytest.mark.parametrize( + 'method,rtol', + [ + ('magnetic_field', 1E-7), + ('magnetic_flux_density', 1E-7), + ] + ) + def test_accuracy(self, method, rtol): + test_prism = getattr(self.prism, method)(self.xyz) + + wx = append_ndim(self.quad_wx, test_prism.ndim) + wy = append_ndim(self.quad_wy, test_prism.ndim) + wz = append_ndim(self.quad_wz, test_prism.ndim) + + test_quad = getattr(self.dipole, method)(self.quad_xyzs) + test_quad *= wx + test_quad *= wy + test_quad *= wz + test_quad = np.sum(test_quad, axis=(0, 1, 2)) + + atol = rtol * (test_prism.max() - test_prism.min()) + npt.assert_allclose(test_quad, test_prism, atol=atol) \ No newline at end of file diff --git a/tests/test_em_tdem_base.py b/tests/em/tdem/test_em_tdem_base.py similarity index 100% rename from tests/test_em_tdem_base.py rename to tests/em/tdem/test_em_tdem_base.py diff --git a/tests/em/tdem/test_em_tdem_electric_dipole.py b/tests/em/tdem/test_em_tdem_electric_dipole.py new file mode 100644 index 00000000..2a8fd641 --- /dev/null +++ b/tests/em/tdem/test_em_tdem_electric_dipole.py @@ -0,0 +1,90 @@ +import numpy as np +import pytest +import numpy.testing as npt + +from geoana.em.tdem import ElectricDipoleWholeSpace + +METHODS = [ + 'vector_potential', + 'magnetic_field', + 'magnetic_field_time_deriv', + 'electric_field', + 'current_density', + 'magnetic_flux_density', + 'magnetic_flux_density_time_deriv', +] + +@pytest.fixture() +def e_dipole(em_dipole_params): + nt = 5 + + current = float(em_dipole_params['current']) + length = float(em_dipole_params['length']) + mu = float(em_dipole_params['mu']) + conductivity = float(em_dipole_params['sigma']) + + ts = np.logspace(-6, -3, nt) + dip = ElectricDipoleWholeSpace( + time=ts, sigma=conductivity, mu=mu, current=current, orientation='x', length=length + ) + return dip + +@pytest.fixture() +def xyz(): + nx, ny, nz = (2, 3, 4) + x = np.linspace(-5, 5, nx) + y = np.linspace(-5, 5, ny) + z = np.linspace(-5, 5, nz) + xyz = np.meshgrid(x, y, z) + return xyz + +@pytest.mark.parametrize('method', METHODS) +@pytest.mark.parametrize('nt', [0, 1, 5]) +def test_broadcasting(e_dipole, xyz, method, nt): + if nt == 0: + e_dipole.time = e_dipole.time[0] + else: + e_dipole.time = e_dipole.time[:nt] + dipole_func = getattr(e_dipole, method) + t_shape = e_dipole.time.shape + + out = dipole_func(xyz) + assert out.shape == (*t_shape, *xyz[0].shape, 3) + + xyz = np.stack(xyz, axis=-1) + out = dipole_func(xyz) + assert out.shape == (*t_shape, *xyz.shape[:-1], 3) + + xyz = xyz.reshape((-1, 3)) + out = dipole_func(xyz) + assert out.shape == (*t_shape, xyz.shape[0], 3) + + xyz = xyz.reshape((-1, 3)) + out = dipole_func(xyz) + assert out.shape == (*t_shape, xyz.shape[0], 3) + + out = dipole_func(xyz[0]) + assert out.shape == (*t_shape, 3) + + +@pytest.mark.parametrize('method', METHODS) +@pytest.mark.parametrize('orient', ['x', 'y', 'z']) +def test_correct(method, orient, e_dipole, xyz, sympy_tdem_ex_dipole): + x, y, z = xyz + e_dipole.orientation = orient + out = getattr(e_dipole, method)(xyz) + + sympy_func = sympy_tdem_ex_dipole[method] + + # cycle the input and output if not x + if orient == 'x': + verify = sympy_func(e_dipole.time[:, None, None, None], x, y, z) + if orient == 'y': + verify = sympy_func(e_dipole.time[:, None, None, None], y, z, x) + verify = verify[..., [2, 0, 1]] + elif orient == 'z': + verify = sympy_func(e_dipole.time[:, None, None, None], z, x, y) + verify = verify[..., [1, 2, 0]] + + # find smallest value that's not zero... + npt.assert_allclose(verify, out, rtol=1E-5) \ No newline at end of file diff --git a/tests/em/tdem/test_em_tdem_halfspace.py b/tests/em/tdem/test_em_tdem_halfspace.py new file mode 100644 index 00000000..a3efb443 --- /dev/null +++ b/tests/em/tdem/test_em_tdem_halfspace.py @@ -0,0 +1,123 @@ +import numpy as np +import numpy.testing as npt + +from geoana.em.tdem import VerticalMagneticDipoleHalfSpace +from geoana.em.tdem.base import theta as theta_func + +from geoana.em.tdem.reference import hp_from_vert_4_72, hz_from_vert_4_69a, dhz_from_vert_4_70, dhp_from_vert_4_74 +from geoana.spatial import cylindrical_to_cartesian, cartesian_to_cylindrical + + +def H_from_Vertical( + XYZ, loc, time, sigma, mu, moment +): + theta = theta_func(time, sigma, mu) + r_vec = XYZ - loc + cyl_locs = cartesian_to_cylindrical(r_vec) + hp = -hp_from_vert_4_72(moment, theta, cyl_locs[..., 0]) + hz = hz_from_vert_4_69a(moment, theta, cyl_locs[..., 0]) + h_cyl = np.stack([hp, np.zeros_like(hp), hz], axis=-1) + return cylindrical_to_cartesian(cyl_locs, h_cyl) + + +def dH_from_Vertical( + XYZ, loc, time, sigma, mu, moment +): + XYZ = np.atleast_2d(XYZ) + + theta = theta_func(time, sigma, mu) + r_vec = XYZ - loc + cyl_locs = cartesian_to_cylindrical(r_vec) + dhp = -dhp_from_vert_4_74(moment, theta, cyl_locs[..., 0], time) + dhz = -dhz_from_vert_4_70(moment, theta, cyl_locs[..., 0], sigma, mu) + h_cyl = np.stack([dhp, np.zeros_like(dhp), dhz], axis=-1) + return cylindrical_to_cartesian(cyl_locs, h_cyl) + + +def B_from_Vertical( + XYZ, loc, time, sigma, mu, moment +): + b = mu * H_from_Vertical(XYZ, loc, time, sigma, mu, moment) + return b + + +def dB_from_Vertical( + XYZ, loc, time, sigma, mu, moment +): + db = mu * dH_from_Vertical(XYZ, loc, time, sigma, mu, moment) + return db + + +class TestVerticalMagneticDipoleHalfSpace: + + def test_magnetic_field(self): + time = 1E-3 + vmdhs = VerticalMagneticDipoleHalfSpace(time=time, moment=1E5) + x = np.linspace(-200., 200., 50) + y = np.linspace(-300., 300., 50) + z = [0, ] + xyz = np.stack(np.meshgrid(x, y, z), axis=-1) + + htest = H_from_Vertical( + xyz, vmdhs.location[:-2], vmdhs.time, vmdhs.sigma, vmdhs.mu, vmdhs.moment + ) + print( + "\n\nTesting Vertical Magnetic Dipole Halfspace Magnetic Field H\n" + ) + + h = vmdhs.magnetic_field(xyz) + npt.assert_allclose(h, htest) + + def test_magnetic_flux(self): + time = 1E-3 + vmdhs = VerticalMagneticDipoleHalfSpace(time=time) + x = np.linspace(-200., 200., 50) + y = np.linspace(-300., 300., 50) + z = [0, ] + xyz = np.stack(np.meshgrid(x, y, z), axis=-1).reshape(-1, 3) + + btest = B_from_Vertical( + xyz, vmdhs.location, vmdhs.time, vmdhs.sigma, vmdhs.mu, vmdhs.moment + ) + print( + "\n\nTesting Vertical Magnetic Dipole Halfspace Magnetic Flux Density B\n" + ) + + b = vmdhs.magnetic_flux_density(xyz) + npt.assert_allclose(b, btest) + + def test_magnetic_field_dt(self): + time = 1E-3 + vmdhs = VerticalMagneticDipoleHalfSpace(time=time) + x = np.linspace(-200., 200., 50) + y = np.linspace(-300., 300., 50) + z = [0, ] + xyz = np.stack(np.meshgrid(x, y, z), axis=-1).reshape(-1, 3) + + dh_test = dH_from_Vertical( + xyz, vmdhs.location, vmdhs.time, vmdhs.sigma, vmdhs.mu, vmdhs.moment + ) + print( + "\n\nTesting Vertical Magnetic Dipole Halfspace Magnetic Field Derivative dH\n" + ) + + dh = vmdhs.magnetic_field_time_derivative(xyz) + npt.assert_allclose(dh_test, dh) + + def test_magnetic_flux_dt(self): + time = 1E-3 + vmdhs = VerticalMagneticDipoleHalfSpace(time=time) + x = np.linspace(-200., 200., 50) + y = np.linspace(-300., 300., 50) + z = [0, ] + xyz = np.stack(np.meshgrid(x, y, z), axis=-1).reshape(-1, 3) + + db_test = dB_from_Vertical( + xyz, vmdhs.location, vmdhs.time, vmdhs.sigma, vmdhs.mu, vmdhs.moment + ) + print( + "\n\nTesting Vertical Magnetic Dipole Halfspace Magnetic Flux Density Derivative dB\n" + ) + + db = vmdhs.magnetic_flux_time_derivative(xyz) + npt.assert_allclose(db_test, db) \ No newline at end of file diff --git a/tests/test_em_tdem_planewave.py b/tests/em/tdem/test_em_tdem_planewave.py similarity index 94% rename from tests/test_em_tdem_planewave.py rename to tests/em/tdem/test_em_tdem_planewave.py index d9807b2d..e5faa609 100644 --- a/tests/test_em_tdem_planewave.py +++ b/tests/em/tdem/test_em_tdem_planewave.py @@ -49,7 +49,7 @@ def test_electric_field(): ey = np.zeros_like(ex) ez = np.zeros_like(ex) - e_vec = tpw.electric_field(xyz)[0] + e_vec = tpw.electric_field(xyz) np.testing.assert_equal(ex, e_vec[:, 0]) np.testing.assert_equal(ey, e_vec[:, 1]) @@ -62,7 +62,7 @@ def test_electric_field(): ex = np.zeros_like(ey) ez = np.zeros_like(ey) - e_vec = tpw.electric_field(xyz)[0] + e_vec = tpw.electric_field(xyz) np.testing.assert_equal(ex, e_vec[:, 0]) np.testing.assert_equal(ey, e_vec[:, 1]) @@ -86,7 +86,7 @@ def test_current_density(): jy = np.zeros_like(jx) jz = np.zeros_like(jx) - j_vec = tpw.current_density(xyz)[0] + j_vec = tpw.current_density(xyz) np.testing.assert_equal(jx, j_vec[:, 0]) np.testing.assert_equal(jy, j_vec[:, 1]) @@ -99,7 +99,7 @@ def test_current_density(): jx = np.zeros_like(jy) jz = np.zeros_like(jy) - j_vec = tpw.current_density(xyz)[0] + j_vec = tpw.current_density(xyz) np.testing.assert_equal(jx, j_vec[:, 0]) np.testing.assert_equal(jy, j_vec[:, 1]) @@ -121,7 +121,7 @@ def test_magnetic_field(): hx = np.zeros_like(hy) hz = np.zeros_like(hy) - h_vec = tpw.magnetic_field(xyz)[0] + h_vec = tpw.magnetic_field(xyz) np.testing.assert_equal(hx, h_vec[..., 0]) np.testing.assert_allclose(hy, h_vec[..., 1], rtol=1E-15) @@ -134,7 +134,7 @@ def test_magnetic_field(): hy = np.zeros_like(hx) hz = np.zeros_like(hx) - h_vec = tpw.magnetic_field(xyz)[0] + h_vec = tpw.magnetic_field(xyz) np.testing.assert_allclose(hx, h_vec[..., 0], rtol=1E-15) np.testing.assert_equal(hy, h_vec[..., 1]) @@ -156,7 +156,7 @@ def test_magnetic_flux_density(): bx = np.zeros_like(by) bz = np.zeros_like(by) - b_vec = tpw.magnetic_flux_density(xyz)[0] + b_vec = tpw.magnetic_flux_density(xyz) np.testing.assert_equal(bx, b_vec[..., 0]) np.testing.assert_allclose(by, b_vec[..., 1], rtol=1E-15) @@ -169,7 +169,7 @@ def test_magnetic_flux_density(): by = np.zeros_like(bx) bz = np.zeros_like(bx) - b_vec = tpw.magnetic_flux_density(xyz)[0] + b_vec = tpw.magnetic_flux_density(xyz) np.testing.assert_allclose(bx, b_vec[..., 0], rtol=1E-15) np.testing.assert_equal(by, b_vec[..., 1]) diff --git a/tests/em/test_em_base.py b/tests/em/test_em_base.py new file mode 100644 index 00000000..12972e28 --- /dev/null +++ b/tests/em/test_em_base.py @@ -0,0 +1,77 @@ +import pytest +from scipy.constants import mu_0, epsilon_0 + +from geoana.em import BaseEM, BaseDipole, BaseElectricDipole, BaseMagneticDipole + +import numpy as np +import numpy.testing as npt + +def test_base_defaults(): + base_em = BaseEM() + assert base_em.sigma == 1 + assert base_em.mu == mu_0 + assert base_em.epsilon == epsilon_0 + +def test_base_errors(): + with pytest.raises(TypeError, match="sigma must be a number"): + BaseEM(sigma='box') + with pytest.raises(ValueError, match="sigma must be greater than 0"): + BaseEM(sigma = -2) + with pytest.raises(TypeError, match="epsilon must be a number.*"): + BaseEM(epsilon = "box") + with pytest.raises(ValueError, match="Invalid epsilon.*"): + BaseEM(epsilon = -2) + with pytest.raises(TypeError, match="mu must be a number.*"): + BaseEM(mu = "box") + with pytest.raises(ValueError, match="mu must be greater than 0"): + BaseEM(mu = -2) + +def test_dipole_defaults(): + base_dipole = BaseDipole() + npt.assert_equal(base_dipole.location, [0, 0, 0]) + npt.assert_equal(base_dipole.orientation, [1, 0, 0]) + + base_dipole.orientation = "Y" + npt.assert_equal(base_dipole.orientation, [0, 1, 0]) + + base_dipole.orientation = "Z" + npt.assert_equal(base_dipole.orientation, [0, 0, 1]) + + base_dipole.orientation = [1, 1, 1] + npt.assert_allclose(base_dipole.orientation, np.sqrt([1/3, 1/3, 1/3])) + +def test_dipole_errors(): + + with pytest.raises(TypeError, match="location must be array_like of float.*"): + BaseDipole(location="box") + with pytest.raises(ValueError, match="location must be array_like with shape.*"): + BaseDipole(location=[0, 0]) + with pytest.raises(ValueError, match="Orientation must be one of {'X','Y','Z'}"): + BaseDipole(orientation="box") + with pytest.raises(TypeError, match="orientation must be str or array_like.*"): + BaseDipole(orientation=["a"]) + with pytest.raises(ValueError, match="orientation must be array_like with shape.*"): + BaseDipole(orientation=[0, 0]) + +def test_e_dipole_defaults(): + base_dipole = BaseElectricDipole() + assert base_dipole.length == 1 + assert base_dipole.current == 1 + +def test_e_dipole_errors(): + + with pytest.raises(TypeError, match="length must be a number, got.*"): + BaseElectricDipole(length="box") + with pytest.raises(ValueError, match="length must be greater than 0"): + BaseElectricDipole(length=0) + with pytest.raises(TypeError, match="current must be a number.*"): + BaseElectricDipole(current="box") + +def test_h_dipole_defaults(): + base_dipole = BaseMagneticDipole() + assert base_dipole.moment == 1 + +def test_h_dipole_errors(): + with pytest.raises(TypeError, match="moment must be a number, got.*"): + BaseMagneticDipole(moment="box") + diff --git a/tests/test_em_simple.py b/tests/em/test_em_simple.py similarity index 100% rename from tests/test_em_simple.py rename to tests/em/test_em_simple.py diff --git a/tests/kernels/test_potential_prism.py b/tests/kernels/test_potential_prism.py new file mode 100644 index 00000000..f6a82b30 --- /dev/null +++ b/tests/kernels/test_potential_prism.py @@ -0,0 +1,164 @@ +import numpy as np +import numpy.testing as npt +import pytest + +import geoana.kernels.potential_field_prism as pf +import geoana.gravity as grav +from geoana.em.static import MagneticPrism +try: + from numba import njit +except ImportError: + njit = None + + +class TestCompiledVsNumpy(): + xyz = np.mgrid[-50:50:51j, -50:50:51j, -50:50:51j] + + def test_assert_using_compiled(self): + assert pf._prism_f is not pf.prism_f + + def test_f(self): + x, y, z = self.xyz + v0 = pf._prism_f(x, y, z) + v1 = pf.prism_f(x, y, z) + npt.assert_allclose(v0, v1) + + def test_fz(self): + x, y, z = self.xyz + v0 = pf._prism_fz(x, y, z) + v1 = pf.prism_fz(x, y, z) + npt.assert_allclose(v0, v1) + + def test_fzz(self): + x, y, z = self.xyz + v0 = pf._prism_fzz(x, y, z) + v1 = pf.prism_fzz(x, y, z) + npt.assert_allclose(v0, v1) + + def test_fzx(self): + x, y, z = self.xyz + v0 = pf._prism_fzx(x, y, z) + v1 = pf.prism_fzx(x, y, z) + npt.assert_allclose(v0, v1) + + def test_fzy(self): + x, y, z = self.xyz + v0 = pf._prism_fzy(x, y, z) + v1 = pf.prism_fzy(x, y, z) + npt.assert_allclose(v0, v1) + + def test_fzzz(self): + x, y, z = self.xyz + v0 = pf._prism_fzzz(x, y, z) + v1 = pf.prism_fzzz(x, y, z) + npt.assert_allclose(v0, v1) + + def test_fxxy(self): + x, y, z = self.xyz + v0 = pf._prism_fxxy(x, y, z) + v1 = pf.prism_fxxy(x, y, z) + npt.assert_allclose(v0, v1) + + def test_fxxz(self): + x, y, z = self.xyz + v0 = pf._prism_fxxz(x, y, z) + v1 = pf.prism_fxxz(x, y, z) + npt.assert_allclose(v0, v1) + + def test_fxyz(self): + x, y, z = self.xyz + v0 = pf._prism_fxyz(x, y, z) + v1 = pf.prism_fxyz(x, y, z) + npt.assert_allclose(v0, v1) + + +PRISM_FUNCTIONS = [ + 'prism_f', + 'prism_fz', + 'prism_fzx', + 'prism_fzy', + 'prism_fzz', + 'prism_fzzz', + 'prism_fxxy', + 'prism_fxxz', + 'prism_fxyz' +] + +@pytest.mark.parametrize('method', PRISM_FUNCTIONS) +def test_prism_correct(method, sympy_potential_prism): + x, y, z = np.mgrid[-10:10:6j, -10:10:4j, -10:10:8j] + dx, dy, dz = 0.1, 0.2, 0.3 + test_func = getattr(pf, method) + verify_func = sympy_potential_prism[method] + verify = 0 + test = 0 + # Prism kernels are accurate for evaluating + # definite integrals so... mimic one here + for ix in [0, 1]: + xp = x + ix * dx + sign_ix = 2 * ix - 1 + for iy in [0, 1]: + yp = y + iy * dy + sign_iy = 2 * iy - 1 + for iz in [0, 1]: + zp = z + iz * dz + sign_iz = 2 * iz - 1 + sign = sign_ix * sign_iy * sign_iz + test = test + sign * test_func(xp, yp, zp) + verify = verify + sign * verify_func(xp, yp, zp) + npt.assert_allclose(test, verify, rtol=1E-6) + + +def test_mag_init_and_errors(): + prism = MagneticPrism([-1, -1, -1], [1, 1, 1]) + np.testing.assert_equal(prism.magnetization, [0.0, 0.0, 1.0]) + + np.testing.assert_equal(prism.moment, [0.0, 0.0, 8.0]) + + with pytest.raises(TypeError): + prism.magnetization = 'abc' + + with pytest.raises(ValueError): + prism.magnetization = 1.0 + + +def test_grav_init_and_errors(): + prism = grav.Prism([-1, -1, -1], [1, 1, 1]) + + assert prism.mass == 8.0 + + with pytest.raises(TypeError): + prism.rho = 'abc' + + +@pytest.mark.skipif(njit is None, reason="numba is not installed.") +@pytest.mark.parametrize('function', [ + pf.prism_f, + pf.prism_fz, + pf.prism_fzx, + pf.prism_fzy, + pf.prism_fzz, + pf.prism_fzzz, + pf.prism_fxxy, + pf.prism_fxxz, + pf.prism_fxyz +]) +def test_numba_jitting_nopython(function): + + # create a vectorized jit function of it: + + x = np.random.rand(10) + y = np.random.rand(10) + z = np.random.rand(10) + @njit + def jitted_func(x, y, z): + n = len(x) + out = np.empty_like(x) + for i in range(n): + out[i] = function(x[i], y[i], z[i]) + return out + + v1 = jitted_func(x, y, z) + v2 = function(x, y, z) + + npt.assert_allclose(v1, v2) diff --git a/tests/test_em_fdem_electric_dipole.py b/tests/test_em_fdem_electric_dipole.py deleted file mode 100644 index 237736e7..00000000 --- a/tests/test_em_fdem_electric_dipole.py +++ /dev/null @@ -1,348 +0,0 @@ -import unittest -import pytest -import numpy as np - -from scipy.constants import mu_0, epsilon_0 -from geoana.em import fdem - -# from SimPEG.EM import FDEM -# from SimPEG import Maps - - -def E_from_EDWS( - XYZ, srcLoc, sig, f, current=1., length=1., orientation='X', kappa=0., - epsr=1., t=0. -): - """E_from_EDWS - Computing the analytic electric fields (E) from an electrical dipole in - a wholespace - - You have the option of computing E for multiple frequencies at a single - reciever location - or a single frequency at multiple locations - - :param numpy.ndarray XYZ: reciever locations at which to evaluate E - :param float epsr: relative permitivitty value (unitless), default is 1.0 - :rtype: numpy.ndarray - :return: Ex, Ey, Ez: arrays containing all 3 components of E evaluated at the specified locations and frequencies. - """ - - mu = mu_0*(1+kappa) - epsilon = epsilon_0*epsr - sig_hat = sig + 1j*fdem.omega(f)*epsilon - - XYZ = np.atleast_2d(XYZ) - - dx = XYZ[:, 0] - srcLoc[0] - dy = XYZ[:, 1] - srcLoc[1] - dz = XYZ[:, 2] - srcLoc[2] - - r = np.sqrt(dx**2. + dy**2. + dz**2.) - # k = np.sqrt( -1j*2.*pi*f*mu*sig ) - k = np.sqrt( - fdem.omega(f)**2. * mu * epsilon - 1j * fdem.omega(f) * mu * sig - ) - - front = current * length / (4.*np.pi*sig_hat * r**3) * np.exp(-1j*k*r) - mid = -k**2 * r**2 + 3*1j*k*r + 3 - - if orientation.upper() == 'X': - Ex = front*((dx**2 / r**2)*mid + (k**2 * r**2 - 1j*k*r-1.)) - Ey = front*(dx*dy / r**2)*mid - Ez = front*(dx*dz / r**2)*mid - return Ex, Ey, Ez - - elif orientation.upper() == 'Y': - # x--> y, y--> z, z-->x - Ey = front*((dy**2 / r**2)*mid + (k**2 * r**2 - 1j*k*r-1.)) - Ez = front*(dy*dz / r**2)*mid - Ex = front*(dy*dx / r**2)*mid - return Ex, Ey, Ez - - elif orientation.upper() == 'Z': - # x --> z, y --> x, z --> y - Ez = front*((dz**2 / r**2)*mid + (k**2 * r**2 - 1j*k*r-1.)) - Ex = front*(dz*dx / r**2)*mid - Ey = front*(dz*dy / r**2)*mid - return Ex, Ey, Ez - - -class TestFDEMdipole(unittest.TestCase): - - def test_defaults(self): - frequency = 1 - edws = fdem.ElectricDipoleWholeSpace(frequency) - assert(edws.sigma == 1) - assert(edws.mu == mu_0) - assert(edws.epsilon == epsilon_0) - assert(np.all(edws.orientation == np.r_[1., 0., 0.])) - assert(edws.length == 1.) - assert(np.all(edws.location == np.r_[0., 0., 0.])) - assert(edws.frequency == 1.) - - def test_errors(self): - edws = fdem.ElectricDipoleWholeSpace(1) - with pytest.raises(TypeError): - edws.current = "box" - with pytest.raises(TypeError): - edws.length = "box" - with pytest.raises(ValueError): - edws.length = -2 - with pytest.raises(TypeError): - edws.sigma = "box" - with pytest.raises(ValueError): - edws.sigma = -2 - with pytest.raises(TypeError): - edws.epsilon = "box" - with pytest.raises(ValueError): - edws.epsilon = -2 - - def compare_fields(name, field, ftest): - - def check_component(name, f, ftest): - geoana_norm = np.linalg.norm(f) - test_norm = np.linalg.norm(ftest) - diff = np.linalg.norm(f-ftest) - passed = np.allclose(f, ftest) - print( - "Testing {} ... geoana: {:1.4e}, compare: {:1.4e}, " - "diff: {:1.4e}, passed?: {}".format( - name, geoana_norm, test_norm, diff, passed - ) - ) - return passed - - passed = [] - for i, orientation in enumerate(['x', 'y', 'z']): - for component in ['real', 'imag']: - passed.append(check_component( - orientation + '_' + component, - getattr(field[:, i], component), - getattr(ftest[:, i], component) - )) - return all(passed) - - def electric_dipole_e(self, orientation): - sigma = np.random.random_integers(1) - frequency = np.random.random_integers(1) - edws = fdem.ElectricDipoleWholeSpace( - orientation=orientation, - sigma=sigma, - frequency=frequency - ) - x = np.linspace(-20., 20., 50) - y = np.linspace(-30., 30., 50) - z = np.linspace(-40., 40., 50) - xyz = np.stack(np.meshgrid(x, y, z), axis=-1).reshape(-1, 3) - - extest, eytest, eztest = E_from_EDWS( - xyz, edws.location, edws.sigma, edws.frequency, - orientation=orientation.upper() - ) - - e = edws.electric_field(xyz) - print( - "\n\nTesting Electric Dipole {} orientation\n".format(orientation) - ) - - passed = self.compare_fields(e, np.vstack([extest, eytest, eztest]).T) - self.assertTrue(passed) - - def test_electric_dipole_x_e(self): - self.electric_dipole_e("x") - - def test_electric_dipole_y_e(self): - self.electric_dipole_e("y") - - def test_electric_dipole_z_e(self): - self.electric_dipole_e("z") - - def test_electric_dipole_tilted_e(self): - - frequency = 1.0 - orientation = np.random.rand(3) - orientation = orientation / np.linalg.norm(orientation) - - edws = fdem.ElectricDipoleWholeSpace( - frequency, orientation=orientation - ) - x = np.linspace(-20., 20., 50) - y = np.linspace(-30., 30., 50) - z = np.linspace(-40., 40., 50) - - xyz = np.stack(np.meshgrid(x, y, z), axis=-1).reshape(-1, 3) - - extest0, eytest0, eztest0 = E_from_EDWS( - xyz, edws.location, edws.sigma, edws.frequency, - orientation='X' - ) - extest1, eytest1, eztest1 = E_from_EDWS( - xyz, edws.location, edws.sigma, edws.frequency, - orientation='Y' - ) - extest2, eytest2, eztest2 = E_from_EDWS( - xyz, edws.location, edws.sigma, edws.frequency, - orientation='Z' - ) - - extest = ( - orientation[0]*extest0 + orientation[1]*extest1 + orientation[2]*extest2 - ) - eytest = ( - orientation[0]*eytest0 + orientation[1]*eytest1 + orientation[2]*eytest2 - ) - eztest = ( - orientation[0]*eztest0 + orientation[1]*eztest1 + orientation[2]*eztest2 - ) - - e = edws.electric_field(xyz) - print( - "\n\nTesting Electric Dipole {} orientation\n".format("45 degree") - ) - - self.compare_fields(e, np.vstack([extest, eytest, eztest]).T) - - -# class TestFDEMdipole_SimPEG(unittest.TestCase): -# -# tol = 1e-1 # error must be an order of magnitude smaller than results -# -# # put the source at the center -# -# def getFaceSrc(self, mesh): -# csx = mesh.hx.min() -# csz = mesh.hz.min() -# srcInd = ( -# (mesh.gridFz[:, 0] < csx) & -# (mesh.gridFz[:, 2] < csz/2.) & (mesh.gridFz[:, 2] > -csz/2.) -# ) -# -# src_vecz = np.zeros(mesh.nFz, dtype=complex) -# src_vecz[srcInd] = 1. -# -# return np.hstack( -# [np.zeros(mesh.vnF[:2].sum(), dtype=complex), src_vecz] -# ) -# -# def getProjections(self, mesh): -# ignore_inside_radius = 5*mesh.hx.min() -# ignore_outside_radius = 40*mesh.hx.min() -# -# def ignoredGridLocs(grid): -# return ( -# ( -# grid[:, 0]**2 + grid[:, 1]**2 + grid[:, 2]**2 < -# ignore_inside_radius**2 -# ) | ( -# grid[:, 0]**2 + grid[:, 1]**2 + grid[:, 2]**2 > -# ignore_outside_radius**2 -# ) -# ) -# -# # Faces -# ignore_me_Fx = ignoredGridLocs(mesh.gridFx) -# ignore_me_Fz = ignoredGridLocs(mesh.gridFz) -# ignore_me = np.hstack([ignore_me_Fx, ignore_me_Fz]) -# keep_me = np.array(~ignore_me, dtype=float) -# Pf = discretize.utils.sdiag(keep_me) -# -# # Edges -# ignore_me_Ey = ignoredGridLocs(mesh.gridEy) -# keep_me_Ey = np.array(~ignore_me_Ey, dtype=float) -# Pe = discretize.utils.sdiag(keep_me_Ey) -# -# return Pf, Pe -# -# def test_e_dipole_v_SimPEG(self): -# -# def compare_w_SimPEG(name, geoana, simpeg): -# -# norm_geoana = np.linalg.norm(geoana) -# norm_simpeg = np.linalg.norm(simpeg) -# diff = np.linalg.norm(geoana - simpeg) -# passed = diff < self.tol * 0.5 * (norm_geoana + norm_simpeg) -# print( -# " {} ... geoana: {:1.4e}, SimPEG: {:1.4e}, diff: {:1.4e}, " -# "passed?: {}".format( -# name, norm_geoana, norm_simpeg, diff, passed -# ) -# ) -# -# return passed -# -# print("\n\nComparing Electric dipole with SimPEG") -# -# sigma_back = 1e-1 -# freqs = np.r_[10., 100.] -# -# csx, ncx, npadx = 0.5, 40, 20 -# ncy = 1 -# csz, ncz, npadz = 0.5, 40, 20 -# -# hx = discretize.utils.meshTensor( -# [(csx, ncx), (csx, npadx, 1.2)] -# ) -# hy = 2*np.pi / ncy * np.ones(ncy) -# hz = discretize.utils.meshTensor( -# [(csz, npadz, -1.2), (csz, ncz), (csz, npadz, 1.2)] -# ) -# -# mesh = discretize.CylMesh([hx, hy, hz], x0='00C') -# -# s_e = self.getFaceSrc(mesh) -# prob = FDEM.Problem3D_h(mesh, sigmaMap=Maps.IdentityMap(mesh)) -# srcList = [FDEM.Src.RawVec_e([], f, s_e) for f in freqs] -# survey = FDEM.Survey(srcList) -# -# prob.pair(survey) -# -# fields = prob.fields(sigma_back*np.ones(mesh.nC)) -# -# length = mesh.hz.min() -# current = np.pi * csx**2 -# -# edws = fdem.ElectricDipoleWholeSpace( -# sigma=sigma_back, length=length, current=current, orientation="z" -# ) -# -# Pf, Pe = self.getProjections(mesh) -# -# j_passed = [] -# h_passed = [] -# for i, f in enumerate(freqs): -# edws.frequency = f -# -# j_xz = [] -# for j, component in zip([0, 2], ['x', 'z']): -# grid = getattr(mesh, "gridF{}".format(component)) -# j_xz.append( -# edws.current_density(grid)[:, j -# ] -# ) -# j_geoana = np.hstack(j_xz) -# h_geoana = edws.magnetic_field(mesh.gridEy)[:, 1] -# -# P_j_geoana = Pf*j_geoana -# P_j_simpeg = Pf*discretize.utils.mkvc(fields[srcList[i], 'j']) -# -# P_h_geoana = Pe*h_geoana -# P_h_simpeg = Pe*discretize.utils.mkvc(fields[srcList[i], 'h']) -# -# print("Testing {} Hz".format(f)) -# -# for comp in ['real', 'imag']: -# j_passed.append(compare_w_SimPEG( -# 'J {}'.format(comp), -# getattr(P_j_geoana, comp), -# getattr(P_j_simpeg, comp) -# )) -# h_passed.append(compare_w_SimPEG( -# 'H {}'.format(comp), -# getattr(P_h_geoana, comp), -# getattr(P_h_simpeg, comp) -# )) -# assert(all(j_passed)) -# assert(all(h_passed)) - - -if __name__ == '__main__': - unittest.main() diff --git a/tests/test_em_fdem_magnetic_dipole.py b/tests/test_em_fdem_magnetic_dipole.py deleted file mode 100644 index 48e8757e..00000000 --- a/tests/test_em_fdem_magnetic_dipole.py +++ /dev/null @@ -1,468 +0,0 @@ -import unittest -import numpy as np -from scipy.constants import mu_0, epsilon_0 -from geoana.em import fdem - - -def H_from_MagneticDipoleWholeSpace( - XYZ, srcLoc, sig, f, current=1., loopArea=1., orientation='X', kappa=0, - epsr=1., t=0. -): - - assert current == 1 - assert loopArea == 1 - assert np.all(srcLoc == np.r_[0., 0., 0.]) - assert kappa == 0 - - mu = mu_0 * (1+kappa) - epsilon = epsilon_0 * epsr - m = current * loopArea - - assert m == 1 - assert mu == mu_0 - assert epsilon == epsilon_0 - - omega = lambda f: 2*np.pi*f - - XYZ = np.atleast_2d(XYZ) - # Check - - dx = XYZ[:, 0]-srcLoc[0] - dy = XYZ[:, 1]-srcLoc[1] - dz = XYZ[:, 2]-srcLoc[2] - - r = np.sqrt(dx**2. + dy**2. + dz**2.) - # k = np.sqrt( -1j*2.*np.pi*f*mu*sig ) - k = np.sqrt(omega(f)**2. * mu * epsilon - 1j*omega(f)*mu*sig) - - front = m / (4. * np.pi * r**3) * np.exp(-1j*k*r) - mid = - k**2 * r**2 + 3*1j*k*r + 3 - - if orientation.upper() == 'X': - Hx = front*((dx**2 / r**2)*mid + (k**2 * r**2 - 1j*k*r - 1.)) - Hy = front*(dx*dy / r**2)*mid - Hz = front*(dx*dz / r**2)*mid - - elif orientation.upper() == 'Y': - # x--> y, y--> z, z-->x - Hy = front * ((dy**2 / r**2)*mid + (k**2 * r**2 - 1j*k*r - 1.)) - Hz = front * (dy*dz / r**2)*mid - Hx = front * (dy*dx / r**2)*mid - - elif orientation.upper() == 'Z': - # x --> z, y --> x, z --> y - Hz = front*((dz**2 / r**2)*mid + (k**2 * r**2 - 1j*k*r - 1.)) - Hx = front*(dz*dx / r**2)*mid - Hy = front*(dz*dy / r**2)*mid - - return Hx, Hy, Hz - - -def B_from_MagneticDipoleWholeSpace( - XYZ, srcLoc, sig, f, current=1., loopArea=1., orientation='X', kappa=0, - epsr=1., t=0. -): - - mu = mu_0 * (1+kappa) - - Hx, Hy, Hz = H_from_MagneticDipoleWholeSpace( - XYZ, srcLoc, sig, f, current=current, loopArea=loopArea, - orientation=orientation, kappa=kappa, epsr=epsr - ) - Bx = mu * Hx - By = mu * Hy - Bz = mu * Hz - return Bx, By, Bz - - -def E_from_MagneticDipoleWholeSpace( - XYZ, srcLoc, sig, f, current=1., loopArea=1., orientation='X', kappa=0., - epsr=1., t=0. -): - - mu = mu_0 * (1+kappa) - epsilon = epsilon_0 * epsr - m = current * loopArea - - omega = lambda f: 2 * np.pi * f - - XYZ = np.atleast_2d(XYZ) - - dx = XYZ[:, 0]-srcLoc[0] - dy = XYZ[:, 1]-srcLoc[1] - dz = XYZ[:, 2]-srcLoc[2] - - r = np.sqrt( dx**2. + dy**2. + dz**2.) - # k = np.sqrt( -1j*2.*np.pi*f*mu*sig ) - k = np.sqrt( omega(f)**2. *mu*epsilon -1j*omega(f)*mu*sig ) - - front = ( - ((1j * omega(f) * mu * m) / (4.* np.pi * r**2)) * - (1j * k * r + 1) * np.exp(-1j*k*r) - ) - - if orientation.upper() == 'X': - Ey = front * (dz / r) - Ez = front * (-dy / r) - Ex = np.zeros_like(Ey) - - elif orientation.upper() == 'Y': - Ex = front * (-dz / r) - Ez = front * (dx / r) - Ey = np.zeros_like(Ex) - - elif orientation.upper() == 'Z': - Ex = front * (dy / r) - Ey = front * (-dx / r) - Ez = np.zeros_like(Ex) - - return Ex, Ey, Ez - - -class TestFDEMdipole(unittest.TestCase): - - def test_defaults(self): - TOL = 1e-15 - frequency = 1 - mdws = fdem.MagneticDipoleWholeSpace(frequency) - assert(mdws.sigma == 1.) - assert(mdws.mu == mu_0) - assert(mdws.epsilon == epsilon_0) - assert(np.all(mdws.orientation == np.r_[1., 0., 0.])) - assert(mdws.moment == 1.) - assert(np.all(mdws.location == np.r_[0., 0., 0.])) - assert(mdws.frequency == 1.) - assert(mdws.omega == 2.*np.pi*1.) - assert(mdws.quasistatic is False) - assert np.linalg.norm( - mdws.wavenumber - np.sqrt( - mu_0 * epsilon_0 * (2*np.pi)**2 - 1j * mu_0 * 1. * 2*np.pi - ) - ) <= TOL - assert np.linalg.norm( - mdws.wavenumber**2 - ( - mu_0 * epsilon_0 * (2*np.pi)**2 - 1j * mu_0 * 1. * 2*np.pi - ) - ) <= TOL - assert np.linalg.norm( - mdws.skin_depth - (np.sqrt( - (mu_0 * epsilon_0 / 2) * (np.sqrt(1 + 1 ** 2 / (2*np.pi * epsilon_0) ** 2) - 1)) / (2*np.pi) - ) - ) <= TOL - - def compare_fields(name, field, ftest): - - def check_component(name, f, ftest): - geoana_norm = np.linalg.norm(f) - test_norm = np.linalg.norm(ftest) - diff = np.linalg.norm(f-ftest) - passed = np.allclose(f, ftest) - print( - "Testing {} ... geoana: {:1.4e}, compare: {:1.4e}, " - "diff: {:1.4e}, passed?: {}".format( - name, geoana_norm, test_norm, diff, passed - ) - ) - return passed - - passed = [] - for i, orientation in enumerate(['x', 'y', 'z']): - for component in ['real', 'imag']: - passed.append(check_component( - orientation + '_' + component, - getattr(field[:, i], component), - getattr(ftest[:, i], component) - )) - return all(passed) - - def magnetic_dipole_b(self, orientation): - sigma = 1 - frequency = 1. - mdws = fdem.MagneticDipoleWholeSpace( - orientation=orientation, - sigma=sigma, - frequency=frequency - ) - x = np.linspace(-20., 20., 50) - y = np.linspace(-30., 30., 50) - z = np.linspace(-40., 40., 50) - xyz = np.stack(np.meshgrid(x, y, z), axis=-1).reshape(-1, 3) - - # srcLoc, obsLoc, component, orientation='Z', moment=1., mu=mu_0 - - # btest = [MagneticDipoleFields( - # mdws.location, xyz, rx_orientation, - # orientation=orientation.upper() - # ) for rx_orientation in ["x", "y", "z"]] - - bxtest, bytest, bztest = B_from_MagneticDipoleWholeSpace( - xyz, mdws.location, mdws.sigma, mdws.frequency, - orientation=orientation - ) - - b = mdws.magnetic_flux_density(xyz) - print( - "\n\nTesting Magnetic Dipole B: {} orientation\n".format(orientation) - ) - - passed = self.compare_fields(b, np.vstack([bxtest, bytest, bztest]).T) - self.assertTrue(passed) - - def magnetic_dipole_e(self, orientation): - sigma = 1e-2 - frequency = 1 - mdws = fdem.MagneticDipoleWholeSpace( - orientation=orientation, - sigma=sigma, - frequency=frequency - ) - x = np.linspace(-20., 20., 50) - y = np.linspace(-30., 30., 50) - z = np.linspace(-40., 40., 50) - xyz = np.stack(np.meshgrid(x, y, z), axis=-1).reshape(-1, 3) - - extest, eytest, eztest = E_from_MagneticDipoleWholeSpace( - xyz, mdws.location, mdws.sigma, mdws.frequency, - orientation=orientation - ) - - e = mdws.electric_field(xyz) - print( - "\n\nTesting Magnetic Dipole E: {} orientation\n".format(orientation) - ) - - passed = self.compare_fields(e, np.vstack([extest, eytest, eztest]).T) - self.assertTrue(passed) - - def test_magnetic_dipole_x_b(self): - self.magnetic_dipole_b("x") - - def test_magnetic_dipole_y_b(self): - self.magnetic_dipole_b("y") - - def test_magnetic_dipole_z_b(self): - self.magnetic_dipole_b("z") - - def test_magnetic_dipole_tilted_b(self): - - frequency = 1.0 - orientation = np.random.rand(3) - orientation = orientation / np.linalg.norm(orientation) - - mdws = fdem.MagneticDipoleWholeSpace( - frequency, orientation=orientation - ) - x = np.linspace(-20., 20., 50) - y = np.linspace(-30., 30., 50) - z = np.linspace(-40., 40., 50) - - xyz = np.stack(np.meshgrid(x, y, z), axis=-1).reshape(-1, 3) - - bxtest0, bytest0, bztest0 = B_from_MagneticDipoleWholeSpace( - xyz, mdws.location, mdws.sigma, mdws.frequency, - orientation='X' - ) - bxtest1, bytest1, bztest1 = B_from_MagneticDipoleWholeSpace( - xyz, mdws.location, mdws.sigma, mdws.frequency, - orientation='Y' - ) - bxtest2, bytest2, bztest2 = B_from_MagneticDipoleWholeSpace( - xyz, mdws.location, mdws.sigma, mdws.frequency, - orientation='Z' - ) - - bxtest = ( - orientation[0]*bxtest0 + orientation[1]*bxtest1 + orientation[2]*bxtest2 - ) - bytest = ( - orientation[0]*bytest0 + orientation[1]*bytest1 + orientation[2]*bytest2 - ) - bztest = ( - orientation[0]*bztest0 + orientation[1]*bztest1 + orientation[2]*bztest2 - ) - - b = mdws.magnetic_flux_density(xyz) - print( - "\n\nTesting Magnetic Dipole B: {} orientation\n".format("45 degree") - ) - - self.compare_fields(b, np.vstack([bxtest, bytest, bztest]).T) - - def test_magnetic_dipole_x_e(self): - self.magnetic_dipole_e("x") - - def test_magnetic_dipole_y_e(self): - self.magnetic_dipole_e("y") - - def test_magnetic_dipole_z_e(self): - self.magnetic_dipole_e("z") - - def test_magnetic_dipole_tilted_e(self): - - frequency = 1.0 - orientation = np.random.rand(3) - orientation = orientation / np.linalg.norm(orientation) - - mdws = fdem.MagneticDipoleWholeSpace( - frequency, orientation=orientation - ) - x = np.linspace(-20., 20., 10) - y = np.linspace(-30., 30., 10) - z = np.linspace(-40., 40., 10) - - xyz = np.stack(np.meshgrid(x, y, z), axis=-1).reshape(-1, 3) - - extest0, eytest0, eztest0 = E_from_MagneticDipoleWholeSpace( - xyz, mdws.location, mdws.sigma, mdws.frequency, - orientation='X' - ) - extest1, eytest1, eztest1 = E_from_MagneticDipoleWholeSpace( - xyz, mdws.location, mdws.sigma, mdws.frequency, - orientation='Y' - ) - extest2, eytest2, eztest2 = E_from_MagneticDipoleWholeSpace( - xyz, mdws.location, mdws.sigma, mdws.frequency, - orientation='Z' - ) - - extest = ( - orientation[0]*extest0 + orientation[1]*extest1 + orientation[2]*extest2 - ) - eytest = ( - orientation[0]*eytest0 + orientation[1]*eytest1 + orientation[2]*eytest2 - ) - eztest = ( - orientation[0]*eztest0 + orientation[1]*eztest1 + orientation[2]*eztest2 - ) - - e = mdws.electric_field(xyz) - print( - "\n\nTesting Magnetic Dipole E: {} orientation\n".format("45 degree") - ) - - self.compare_fields(e, np.vstack([extest, eytest, eztest]).T) - - -# class TestFDEMdipole_SimPEG(unittest.TestCase): -# -# tol = 1e-1 # error must be an order of magnitude smaller than results -# -# def getProjections(self, mesh): -# ignore_inside_radius = 10*mesh.hx.min() -# ignore_outside_radius = 40*mesh.hx.min() -# -# def ignoredGridLocs(grid): -# return ( -# ( -# grid[:, 0]**2 + grid[:, 1]**2 + grid[:, 2]**2 < -# ignore_inside_radius**2 -# ) | ( -# grid[:, 0]**2 + grid[:, 1]**2 + grid[:, 2]**2 > -# ignore_outside_radius**2 -# ) -# ) -# -# # Faces -# ignore_me_Fx = ignoredGridLocs(mesh.gridFx) -# ignore_me_Fz = ignoredGridLocs(mesh.gridFz) -# ignore_me = np.hstack([ignore_me_Fx, ignore_me_Fz]) -# keep_me = np.array(~ignore_me, dtype=float) -# Pf = discretize.utils.sdiag(keep_me) -# -# # Edges -# ignore_me_Ey = ignoredGridLocs(mesh.gridEy) -# keep_me_Ey = np.array(~ignore_me_Ey, dtype=float) -# Pe = discretize.utils.sdiag(keep_me_Ey) -# -# return Pf, Pe -# -# def test_b_dipole_v_SimPEG(self): -# -# def compare_w_SimPEG(name, geoana, simpeg): -# -# norm_geoana = np.linalg.norm(geoana) -# norm_simpeg = np.linalg.norm(simpeg) -# diff = np.linalg.norm(geoana - simpeg) -# passed = diff < self.tol * 0.5 * (norm_geoana + norm_simpeg) -# print( -# " {} ... geoana: {:1.4e}, SimPEG: {:1.4e}, diff: {:1.4e}, " -# "passed?: {}".format( -# name, norm_geoana, norm_simpeg, diff, passed -# ) -# ) -# -# return passed -# -# print("\n\nComparing Magnetic dipole with SimPEG") -# -# sigma_back = 1. -# freqs = np.r_[10., 100.] -# -# csx, ncx, npadx = 1, 50, 50 -# ncy = 1 -# csz, ncz, npadz = 1, 50, 50 -# -# hx = discretize.utils.meshTensor( -# [(csx, ncx), (csx, npadx, 1.3)] -# ) -# hy = 2*np.pi / ncy * np.ones(ncy) -# hz = discretize.utils.meshTensor( -# [(csz, npadz, -1.3), (csz, ncz), (csz, npadz, 1.3)] -# ) -# -# mesh = discretize.CylMesh([hx, hy, hz], x0='00C') -# -# prob = FDEM.Problem3D_e(mesh, sigmaMap=Maps.IdentityMap(mesh)) -# srcList = [FDEM.Src.MagDipole([], loc=np.r_[0., 0., 0.], freq=f) for f in freqs] -# survey = FDEM.Survey(srcList) -# -# prob.pair(survey) -# -# fields = prob.fields(sigma_back*np.ones(mesh.nC)) -# -# moment = 1. -# mdws = fdem.MagneticDipoleWholeSpace( -# sigma=sigma_back, moment=moment, orientation="z" -# ) -# -# Pf, Pe = self.getProjections(mesh) -# -# e_passed = [] -# b_passed = [] -# for i, f in enumerate(freqs): -# mdws.frequency = f -# -# b_xz = [] -# for b, component in zip([0, 2], ['x', 'z']): -# grid = getattr(mesh, "gridF{}".format(component)) -# b_xz.append( -# mdws.magnetic_flux_density(grid)[:, b] -# ) -# b_geoana = np.hstack(b_xz) -# e_geoana = mdws.electric_field(mesh.gridEy)[:, 1] -# -# P_e_geoana = Pe*e_geoana -# P_e_simpeg = Pe*discretize.utils.mkvc(fields[srcList[i], 'e']) -# -# P_b_geoana = Pf*b_geoana -# P_b_simpeg = Pf*discretize.utils.mkvc(fields[srcList[i], 'b']) -# -# print("Testing {} Hz".format(f)) -# -# for comp in ['real', 'imag']: -# e_passed.append(compare_w_SimPEG( -# 'E {}'.format(comp), -# getattr(P_e_geoana, comp), -# getattr(P_e_simpeg, comp) -# )) -# b_passed.append(compare_w_SimPEG( -# 'B {}'.format(comp), -# getattr(P_b_geoana, comp), -# getattr(P_b_simpeg, comp) -# )) -# assert(all(e_passed)) -# assert(all(b_passed)) - - -if __name__ == '__main__': - unittest.main() diff --git a/tests/test_em_tdem_halfspace.py b/tests/test_em_tdem_halfspace.py deleted file mode 100644 index 4ee3516b..00000000 --- a/tests/test_em_tdem_halfspace.py +++ /dev/null @@ -1,150 +0,0 @@ -import numpy as np - -from geoana.em.tdem import VerticalMagneticDipoleHalfSpace -from geoana.em.tdem.base import theta -from scipy.special import erf, ive - - -def H_from_Vertical( - XYZ, loc, time, sigma, mu, moment -): - - XYZ = np.atleast_2d(XYZ) - - r_vec = XYZ - loc - r = np.linalg.norm(r_vec[:, :2], axis=-1) - x = r_vec[:, 0] - y = r_vec[:, 1] - thr = theta(time, sigma, mu=mu)[:, None] * r - - h_z = 1.0 / r ** 3 * ( - (9 / (2 * thr ** 2) - 1) * erf(thr) - - (9 / thr + 4 * thr) / np.sqrt(np.pi) * np.exp(-thr ** 2) - ) - - h_r = 2 * thr ** 2 / r ** 3 * ( - ive(1, thr ** 2 / 2) - ive(2, thr ** 2 / 2) - ) - - angle = np.arctan2(y, x) - h_x = np.cos(angle) * h_r - h_y = np.sin(angle) * h_r - - h = moment / (4 * np.pi) * np.stack((h_x, h_y, h_z), axis=-1) - return h[0] - - -def dH_from_Vertical( - XYZ, loc, time, sigma, mu, moment -): - XYZ = np.atleast_2d(XYZ) - - r_vec = XYZ - loc - r = np.linalg.norm(r_vec[:, :2], axis=-1) - x = r_vec[:, 0] - y = r_vec[:, 1] - tr = theta(time, sigma, mu)[:, None] * r - - dhz_dt = 1 / (r ** 3 * time[:, None]) * ( - 9 / (2 * tr ** 2) * erf(tr) - - (4 * tr ** 3 + 6 * tr + 9 / tr) / np.sqrt(np.pi) * np.exp(-tr ** 2) - ) - - dhr_dt = - 2 * tr ** 2 / (r ** 3 * time[:, None]) * ( - (1 + tr ** 2) * ive(0, tr ** 2 / 2) - - (2 + tr ** 2 + 4 / tr ** 2) * ive(1, tr ** 2 / 2) - ) - angle = np.arctan2(y, x) - dhx_dt = np.cos(angle) * dhr_dt - dhy_dt = np.sin(angle) * dhr_dt - dh_dt = moment / (4 * np.pi) * np.stack((dhx_dt, dhy_dt, dhz_dt), axis=-1) - return dh_dt[0] - - -def B_from_Vertical( - XYZ, loc, time, sigma, mu, moment -): - b = mu * H_from_Vertical(XYZ, loc, time, sigma, mu, moment) - return b - - -def dB_from_Vertical( - XYZ, loc, time, sigma, mu, moment -): - db = mu * dH_from_Vertical(XYZ, loc, time, sigma, mu, moment) - return db - - -class TestVerticalMagneticDipoleHalfSpace: - - def test_magnetic_field(self): - time = 1 - vmdhs = VerticalMagneticDipoleHalfSpace(time=time) - x = np.linspace(-20., 20., 50) - y = np.linspace(-30., 30., 50) - z = np.linspace(-40., 40., 50) - xyz = np.stack(np.meshgrid(x, y, z), axis=-1).reshape(-1, 3) - - htest = H_from_Vertical( - xyz, vmdhs.location, vmdhs.time, vmdhs.sigma, vmdhs.mu, vmdhs.moment - ) - print( - "\n\nTesting Vertical Magnetic Dipole Halfspace Magnetic Field H\n" - ) - - h = vmdhs.magnetic_field(xyz) - np.testing.assert_equal(htest, h) - - def test_magnetic_flux(self): - time = 1 - vmdhs = VerticalMagneticDipoleHalfSpace(time=time) - x = np.linspace(-20., 20., 50) - y = np.linspace(-30., 30., 50) - z = np.linspace(-40., 40., 50) - xyz = np.stack(np.meshgrid(x, y, z), axis=-1).reshape(-1, 3) - - btest = B_from_Vertical( - xyz, vmdhs.location, vmdhs.time, vmdhs.sigma, vmdhs.mu, vmdhs.moment - ) - print( - "\n\nTesting Vertical Magnetic Dipole Halfspace Magnetic Flux Density B\n" - ) - - b = vmdhs.magnetic_flux_density(xyz) - np.testing.assert_equal(btest, b) - - def test_magnetic_field_dt(self): - time = 1 - vmdhs = VerticalMagneticDipoleHalfSpace(time=time) - x = np.linspace(-20., 20., 50) - y = np.linspace(-30., 30., 50) - z = np.linspace(-40., 40., 50) - xyz = np.stack(np.meshgrid(x, y, z), axis=-1).reshape(-1, 3) - - dh_test = dH_from_Vertical( - xyz, vmdhs.location, vmdhs.time, vmdhs.sigma, vmdhs.mu, vmdhs.moment - ) - print( - "\n\nTesting Vertical Magnetic Dipole Halfspace Magnetic Field Derivative dH\n" - ) - - dh = vmdhs.magnetic_field_time_derivative(xyz) - np.testing.assert_equal(dh_test, dh) - - def test_magnetic_flux_dt(self): - time = 1 - vmdhs = VerticalMagneticDipoleHalfSpace(time=time) - x = np.linspace(-20., 20., 50) - y = np.linspace(-30., 30., 50) - z = np.linspace(-40., 40., 50) - xyz = np.stack(np.meshgrid(x, y, z), axis=-1).reshape(-1, 3) - - db_test = dB_from_Vertical( - xyz, vmdhs.location, vmdhs.time, vmdhs.sigma, vmdhs.mu, vmdhs.moment - ) - print( - "\n\nTesting Vertical Magnetic Dipole Halfspace Magnetic Flux Density Derivative dB\n" - ) - - db = vmdhs.magnetic_flux_time_derivative(xyz) - np.testing.assert_equal(db_test, db) \ No newline at end of file diff --git a/tests/test_gravity.py b/tests/test_gravity.py index 72d6bf94..6ddf0ecb 100644 --- a/tests/test_gravity.py +++ b/tests/test_gravity.py @@ -1,49 +1,16 @@ import pytest import numpy as np -from scipy.constants import G +import numpy.testing as npt +from scipy.special import roots_legendre from geoana import gravity +from geoana.utils import append_ndim - -def U_from_PointMass( - XYZ, loc, m -): - - XYZ = np.atleast_2d(XYZ) - - r_vec = XYZ - loc - r = np.linalg.norm(r_vec, axis=-1) - - u_g = (G * m) / r - return u_g - - -def g_from_PointMass( - XYZ, loc, m -): - - XYZ = np.atleast_2d(XYZ) - - r_vec = XYZ - loc - r = np.linalg.norm(r_vec, axis=-1) - - g_vec = -G * m * r_vec / r[..., None] ** 3 - return g_vec - - -def gtens_from_PointMass( - XYZ, loc, m -): - - XYZ = np.atleast_2d(XYZ) - - r_vec = XYZ - loc - r = np.linalg.norm(r_vec, axis=-1) - - g_tens = -G * m * (np.eye(3) / r[..., None, None] ** 3 - - 3 * r_vec[..., None] * r_vec[..., None, :] / r[..., None, None] ** 5) - return g_tens - +METHODS = [ + "gravitational_potential", + "gravitational_field", + "gravitational_gradient", +] class TestPointMass: @@ -63,119 +30,24 @@ def test_errors(self): with pytest.raises(TypeError): pm.location = ["string"] - def test_gravitational_potential(self): - mass = 1.0 - pm = gravity.PointMass( + @pytest.mark.parametrize("method", METHODS) + def test_correct(self, method, sympy_grav_point, grav_point_params): + mass = grav_point_params["mass"] + grav_obj = gravity.PointMass( mass=mass ) x = np.linspace(-20., 20., 50) y = np.linspace(-30., 30., 50) z = np.linspace(-40., 40., 50) - xyz = np.stack(np.meshgrid(x, y, z), axis=-1).reshape(-1, 3) - - utest = U_from_PointMass( - xyz, pm.location, pm.mass - ) - - u = pm.gravitational_potential(xyz) - print( - "\n\nTesting Gravitational Potential U for Point Mass\n" - ) - - np.testing.assert_equal(utest, u) - - def test_gravitational_field(self): - mass = 1.0 - pm = gravity.PointMass( - mass=mass - ) - x = np.linspace(-20., 20., 50) - y = np.linspace(-30., 30., 50) - z = np.linspace(-40., 40., 50) - xyz = np.stack(np.meshgrid(x, y, z), axis=-1).reshape(-1, 3) - - gtest = g_from_PointMass( - xyz, pm.location, pm.mass - ) - - g = pm.gravitational_field(xyz) - print( - "\n\nTesting Gravitational Field g for Point Mass\n" - ) - - np.testing.assert_equal(gtest, g) - - def test_gravitational_gradient(self): - mass = 1.0 - pm = gravity.PointMass( - mass=mass - ) - x = np.linspace(-20., 20., 5) - y = np.linspace(-30., 30., 5) - z = np.linspace(-40., 40., 5) - xyz = np.stack(np.meshgrid(x, y, z), axis=-1).reshape(-1, 3) - - g_tenstest = gtens_from_PointMass( - xyz, pm.location, pm.mass - ) - - g_tens = pm.gravitational_gradient(xyz) - print( - "\n\nTesting Gravitational Gradient g_tens for Point Mass\n" - ) - - np.testing.assert_equal(g_tenstest, g_tens) - - -def U_from_Sphere( - XYZ, loc, m, rho, radius -): + xyz = np.meshgrid(x, y, z) - XYZ = np.atleast_2d(XYZ) + func = getattr(grav_obj, method) - r_vec = XYZ - loc - r = np.linalg.norm(r_vec, axis=-1) + out = func(xyz) - u_g = np.zeros_like(r) - ind0 = r > radius - u_g[ind0] = (G * m) / r[ind0] - u_g[~ind0] = G * 2 / 3 * np.pi * rho * (3 * radius ** 2 - r[~ind0] ** 2) - return u_g + verify = sympy_grav_point[method](*xyz) - -def g_from_Sphere( - XYZ, loc, m, rho, radius -): - - XYZ = np.atleast_2d(XYZ) - - r_vec = XYZ - loc - r = np.linalg.norm(r_vec, axis=-1) - - g_vec = np.zeros((*r.shape, 3)) - ind0 = r > radius - g_vec[ind0] = -G * m * r_vec[ind0] / r[ind0, None] ** 3 - g_vec[~ind0] = -G * 4 / 3 * np.pi * rho * r_vec[~ind0] - return g_vec - - -def gtens_from_Sphere( - XYZ, loc, m, rho, radius -): - - XYZ = np.atleast_2d(XYZ) - - r_vec = XYZ - loc - r = np.linalg.norm(r_vec, axis=-1) - - g_tens = np.zeros((*r.shape, 3, 3)) - ind0 = r > radius - ind1 = r == radius - g_tens[ind0] = -G * m * (np.eye(3) / r[ind0, None, None] ** 3 - - 3 * r_vec[ind0, None] * r_vec[ind0, None, :] / r[ind0, None, None] ** 5) - g_tens[~ind0] = -G * 4 / 3 * np.pi * rho * np.eye(3) - g_tens[ind1] = np.nan - return g_tens + np.testing.assert_allclose(out, verify) class TestSphere: @@ -202,75 +74,70 @@ def test_errors(self): with pytest.raises(TypeError): s.location = ["string"] - def test_gravitational_potential(self): - radius = 1.0 - rho = 1.0 - location = [0., 0., 0.] - s = gravity.Sphere( + @pytest.mark.parametrize("method", METHODS) + def test_correct(self, method, sympy_grav_sphere, grav_point_params): + mass = grav_point_params["mass"] + radius = grav_point_params["radius"] + vol = 4/3 * np.pi * radius**3 + rho = mass / vol + + grav_obj = gravity.Sphere( radius=radius, rho=rho, - location=location ) - x = np.linspace(-20., 20., 50) - y = np.linspace(-30., 30., 50) - z = np.linspace(-40., 40., 50) - xyz = np.stack(np.meshgrid(x, y, z), axis=-1).reshape(-1, 3) + x = np.linspace(-2., 2., 50) + y = np.linspace(-3., 3., 50) + z = np.linspace(-4., 4., 50) + xyz = np.meshgrid(x, y, z) - utest = U_from_Sphere( - xyz, s.location, s.mass, s.rho, s.radius - ) - print( - "\n\nTesting Gravitational Potential U for Sphere\n" - ) + func = getattr(grav_obj, method) - u = s.gravitational_potential(xyz) - np.testing.assert_equal(utest, u) + out = func(xyz) - def test_gravitational_field(self): - radius = 1.0 - rho = 1.0 - location = [0., 0., 0.] - s = gravity.Sphere( - radius=radius, - rho=rho, - location=location - ) - x = np.linspace(-20., 20., 50) - y = np.linspace(-30., 30., 50) - z = np.linspace(-40., 40., 50) - xyz = np.stack(np.meshgrid(x, y, z), axis=-1).reshape(-1, 3) + verify = sympy_grav_sphere[method](*xyz) - gtest = g_from_Sphere( - xyz, s.location, s.mass, s.rho, s.radius - ) - print( - "\n\nTesting Gravitational Field g for Sphere\n" - ) + np.testing.assert_allclose(out, verify) - g = s.gravitational_field(xyz) - np.testing.assert_equal(gtest, g) - def test_gravitational_gradient(self): - radius = 1.0 - rho = 1.0 - location = [0., 0., 0.] - s = gravity.Sphere( - radius=radius, - rho=rho, - location=location - ) - x = np.linspace(-20., 20., 50) - y = np.linspace(-30., 30., 50) - z = np.linspace(-40., 40., 50) - xyz = np.stack(np.meshgrid(x, y, z), axis=-1).reshape(-1, 3) +class TestGravityAccuracy(): + x, y, z = np.mgrid[-100:100:20j, -100:100:20j, -100:100:20j] + xyz = np.stack((x, y, z), axis=-1).reshape((-1, 3)) + dx = 0.1 + prism = gravity.Prism(dx * np.r_[-1, -1, -1], dx * np.r_[1, 1, 1], rho=2) - g_tens_test = gtens_from_Sphere( - xyz, s.location, s.mass, s.rho, s.radius - ) - print( - "\n\nTesting Gravitational Gradient g_tens for Sphere\n" - ) + quad_points, quad_weights = roots_legendre(5) + quad_points = (prism.max_location - prism.min_location)[:, None] * (quad_points + 1) / 2 + prism.min_location[:, + None] + quad_points = np.stack(np.meshgrid(*quad_points, indexing='ij'), axis=-1) + quad_wx, quad_wy, quad_wz = quad_weights * (prism.max_location - prism.min_location)[:, None] / 2 + quad_wx = quad_wx[:, None, None] + quad_wy = quad_wy[None, :, None] + quad_wz = quad_wz[None, None, :] + + pm = gravity.PointMass(mass=prism.rho, location=[0, 0, 0]) + + quad_xyzs = xyz - quad_points[..., None, :] + + @pytest.mark.parametrize( + 'method,rtol', + [ + ('gravitational_potential', 1E-7), + ('gravitational_field', 1E-7), + ('gravitational_gradient', 1E-7), + ] + ) + def test_accuracy(self, method, rtol): + test_prism = getattr(self.prism, method)(self.xyz) + + wx = append_ndim(self.quad_wx, test_prism.ndim) + wy = append_ndim(self.quad_wy, test_prism.ndim) + wz = append_ndim(self.quad_wz, test_prism.ndim) - g_tens = s.gravitational_gradient(xyz) - np.testing.assert_allclose(g_tens_test, g_tens, atol=1E-9) + test_quad = getattr(self.pm, method)(self.quad_xyzs) + test_quad *= wx + test_quad *= wy + test_quad *= wz + test_quad = np.sum(test_quad, axis=(0, 1, 2)) + atol = rtol * (test_prism.max() - test_prism.min()) + npt.assert_allclose(test_quad, test_prism, atol=atol) \ No newline at end of file diff --git a/tests/test_potential_prism.py b/tests/test_potential_prism.py deleted file mode 100644 index 01097725..00000000 --- a/tests/test_potential_prism.py +++ /dev/null @@ -1,327 +0,0 @@ -import numpy as np -from numpy.testing import assert_allclose -import pytest - -import geoana.kernels.potential_field_prism as pf -import geoana.gravity as grav -from geoana.em.static import MagneticPrism, MagneticDipoleWholeSpace -try: - from numba import njit -except ImportError: - njit = None - -try: - from discretize.tests import check_derivative -except ImportError: - check_derivative = None - - -class TestCompiledVsNumpy(): - xyz = np.mgrid[-50:50:51j, -50:50:51j, -50:50:51j] - - def test_assert_using_compiled(self): - assert pf._prism_f is not pf.prism_f - - def test_f(self): - x, y, z = self.xyz - v0 = pf._prism_f(x, y, z) - v1 = pf.prism_f(x, y, z) - assert_allclose(v0, v1) - - def test_fz(self): - x, y, z = self.xyz - v0 = pf._prism_fz(x, y, z) - v1 = pf.prism_fz(x, y, z) - assert_allclose(v0, v1) - - def test_fzz(self): - x, y, z = self.xyz - v0 = pf._prism_fzz(x, y, z) - v1 = pf.prism_fzz(x, y, z) - assert_allclose(v0, v1) - - def test_fzx(self): - x, y, z = self.xyz - v0 = pf._prism_fzx(x, y, z) - v1 = pf.prism_fzx(x, y, z) - assert_allclose(v0, v1) - - def test_fzy(self): - x, y, z = self.xyz - v0 = pf._prism_fzy(x, y, z) - v1 = pf.prism_fzy(x, y, z) - assert_allclose(v0, v1) - - def test_fzzz(self): - x, y, z = self.xyz - v0 = pf._prism_fzzz(x, y, z) - v1 = pf.prism_fzzz(x, y, z) - assert_allclose(v0, v1) - - def test_fxxy(self): - x, y, z = self.xyz - v0 = pf._prism_fxxy(x, y, z) - v1 = pf.prism_fxxy(x, y, z) - assert_allclose(v0, v1) - - def test_fxxz(self): - x, y, z = self.xyz - v0 = pf._prism_fxxz(x, y, z) - v1 = pf.prism_fxxz(x, y, z) - assert_allclose(v0, v1) - - def test_fxyz(self): - x, y, z = self.xyz - v0 = pf._prism_fxyz(x, y, z) - v1 = pf.prism_fxyz(x, y, z) - assert_allclose(v0, v1) - - -@pytest.mark.skipif(check_derivative is None, reason="discretize not installed.") -class TestGravityPrismDerivatives(): - xyz = np.mgrid[-100:100:26j, -100:100:26j, -100:100:26j] - x = xyz[0].ravel() - y = xyz[1].ravel() - z = xyz[2].ravel() - prism = grav.Prism([-10, -10, -10], [10, 10, 10], rho=2.5) - - def test_pot_and_field_x(self): - def grav_pot(x): - pot = self.prism.gravitational_potential((x, self.y, self.z)) - - Gv = self.prism.gravitational_field((x, self.y, self.z)) - def J(v): - return Gv[..., 0] * v - return pot, J - - assert check_derivative(grav_pot, self.x, num=3, plotIt=False) - - def test_pot_and_field_y(self): - def grav_pot(y): - pot = self.prism.gravitational_potential((self.x, y, self.z)) - - Gv = self.prism.gravitational_field((self.x, y, self.z)) - def J(v): - return Gv[..., 1] * v - return pot, J - - assert check_derivative(grav_pot, self.y, num=3, plotIt=False) - - def test_pot_and_field_z(self): - def grav_pot(z): - pot = self.prism.gravitational_potential((self.x, self.y, z)) - - Gv = self.prism.gravitational_field((self.x, self.y, z)) - def J(v): - return Gv[..., 2] * v - return pot, J - - assert check_derivative(grav_pot, self.z, num=3, plotIt=False) - - def test_field_and_grad_x(self): - def grav_field(x): - Gv = self.prism.gravitational_field((x, self.y, self.z)).ravel() - GG_x = self.prism.gravitational_gradient((x, self.y, self.z))[..., 0, :] - def J(v): - return (GG_x * v[:, None]).ravel() - return Gv, J - - assert check_derivative(grav_field, self.x, num=3, plotIt=False) - - def test_field_and_grad_y(self): - def grav_field(y): - Gv = self.prism.gravitational_field((self.x, y, self.z)).ravel() - GG_y = self.prism.gravitational_gradient((self.x, y, self.z))[..., 1, :] - def J(v): - return (GG_y * v[:, None]).ravel() - return Gv, J - - assert check_derivative(grav_field, self.y, num=3, plotIt=False) - - def test_field_and_grad_z(self): - def grav_field(z): - Gv = self.prism.gravitational_field((self.x, self.y, z)).ravel() - GG_z = self.prism.gravitational_gradient((self.x, self.y, z))[..., 2, :] - def J(v): - return (GG_z * v[:, None]).ravel() - return Gv, J - - assert check_derivative(grav_field, self.z, num=3, plotIt=False) - - -@pytest.mark.skipif(check_derivative is None, reason="discretize not installed.") -class TestMagPrismDerivatives(): - xyz = np.mgrid[-100:100:26j, -100:100:26j, -100:100:26j] - x = xyz[0].ravel() - y = xyz[1].ravel() - z = xyz[2].ravel() - prism = MagneticPrism([-10, -10, -10], [10, 10, 10], magnetization=[-1, 2, 3]) - - def test_pot_and_field_x(self): - def grav_pot(x): - pot = self.prism.scalar_potential((x, self.y, self.z)) - - Gv = self.prism.magnetic_field((x, self.y, self.z)) - def J(v): - return Gv[..., 0] * v - return pot, J - - assert check_derivative(grav_pot, self.x, num=3, plotIt=False) - - def test_pot_and_field_y(self): - def grav_pot(y): - pot = self.prism.scalar_potential((self.x, y, self.z)) - - Gv = self.prism.magnetic_field((self.x, y, self.z)) - def J(v): - return Gv[..., 1] * v - return pot, J - - assert check_derivative(grav_pot, self.y, num=3, plotIt=False) - - def test_pot_and_field_z(self): - def grav_pot(z): - pot = self.prism.scalar_potential((self.x, self.y, z)) - - Gv = self.prism.magnetic_field((self.x, self.y, z)) - def J(v): - return Gv[..., 2] * v - return pot, J - - assert check_derivative(grav_pot, self.z, num=3, plotIt=False) - - def test_field_and_grad_x(self): - def grav_field(x): - Gv = self.prism.magnetic_field((x, self.y, self.z)).ravel() - GG_x = self.prism.magnetic_field_gradient((x, self.y, self.z))[..., 0, :] - def J(v): - return (GG_x * v[:, None]).ravel() - return Gv, J - - assert check_derivative(grav_field, self.x, num=3, plotIt=False) - - def test_field_and_grad_y(self): - def grav_field(y): - Gv = self.prism.magnetic_field((self.x, y, self.z)).ravel() - GG_y = self.prism.magnetic_field_gradient((self.x, y, self.z))[..., 1, :] - def J(v): - return (GG_y * v[:, None]).ravel() - return Gv, J - - assert check_derivative(grav_field, self.y, num=3, plotIt=False) - - def test_field_and_grad_z(self): - def grav_field(z): - Gv = self.prism.magnetic_field((self.x, self.y, z)).ravel() - GG_z = self.prism.magnetic_field_gradient((self.x, self.y, z))[..., 2, :] - def J(v): - return (GG_z * v[:, None]).ravel() - return Gv, J - - assert check_derivative(grav_field, self.z, num=3, plotIt=False) - - -class TestGravityAccuracy(): - x, y, z = np.mgrid[-100:100:101j, -100:100:101j, -100:100:101j] - xyz = np.stack((x, y, z), axis=-1) - xyz = xyz[np.linalg.norm(xyz, axis=-1) >= 10] - prism = grav.Prism([-1, -1, -1], [1, 1, 1], 2.0) - point = grav.PointMass(mass=prism.mass, location=prism.location) - - def test_potential(self): - pot_prism = self.prism.gravitational_potential(self.xyz) - pot_point = self.point.gravitational_potential(self.xyz) - np.testing.assert_allclose(pot_prism, pot_point, rtol=1E-4) - - def test_field(self): - gv_prism = self.prism.gravitational_field(self.xyz) - gv_point = self.point.gravitational_field(self.xyz) - np.testing.assert_allclose(gv_prism, gv_point, rtol=1E-3) - - def test_gradient(self): - gg_prism = self.prism.gravitational_field(self.xyz) - gg_point = self.point.gravitational_field(self.xyz) - np.testing.assert_allclose( - gg_prism, gg_point, atol=(gg_prism.max() - gg_prism.min())*1E-3 - ) - - -class TestMagneticAccuracy(): - x, y, z = np.mgrid[-100:100:101j, -100:100:101j, -100:100:101j] - xyz = np.stack((x, y, z), axis=-1) - xyz = xyz[np.linalg.norm(xyz, axis=-1) >= 10] - prism = MagneticPrism([-1, -1, -1], [1, 1, 1], magnetization=[-1, 2, -0.5]) - m_mag = np.linalg.norm(prism.moment) - m_unit = prism.moment/m_mag - dipole = MagneticDipoleWholeSpace( - location=prism.location, moment=m_mag, orientation=m_unit - ) - - def test_mag_field(self): - H_prism = self.prism.magnetic_field(self.xyz) - H_dipole = self.dipole.magnetic_field(self.xyz) - np.testing.assert_allclose( - H_prism, H_dipole, atol=(H_prism.max()-H_prism.min())*1E-3 - ) - - def test_mag_flux(self): - B_prism = self.prism.magnetic_flux_density(self.xyz) - B_dipole = self.dipole.magnetic_flux_density(self.xyz) - np.testing.assert_allclose( - B_prism, B_dipole, atol=(B_prism.max()-B_dipole.min())*1E-3 - ) - - -def test_mag_init_and_errors(): - prism = MagneticPrism([-1, -1, -1], [1, 1, 1]) - np.testing.assert_equal(prism.magnetization, [0.0, 0.0, 1.0]) - - np.testing.assert_equal(prism.moment, [0.0, 0.0, 8.0]) - - with pytest.raises(TypeError): - prism.magnetization = 'abc' - - with pytest.raises(ValueError): - prism.magnetization = 1.0 - - -def test_grav_init_and_errors(): - prism = grav.Prism([-1, -1, -1], [1, 1, 1]) - - assert prism.mass == 8.0 - - with pytest.raises(TypeError): - prism.rho = 'abc' - - -@pytest.mark.skipif(njit is None, reason="numba is not installed.") -@pytest.mark.parametrize('function', [ - pf.prism_f, - pf.prism_fz, - pf.prism_fzx, - pf.prism_fzy, - pf.prism_fzz, - pf.prism_fzzz, - pf.prism_fxxy, - pf.prism_fxxz, - pf.prism_fxyz -]) -def test_numba_jitting_nopython(function): - - # create a vectorized jit function of it: - - x = np.random.rand(10) - y = np.random.rand(10) - z = np.random.rand(10) - @njit - def jitted_func(x, y, z): - n = len(x) - out = np.empty_like(x) - for i in range(n): - out[i] = function(x[i], y[i], z[i]) - return out - - v1 = jitted_func(x, y, z) - v2 = function(x, y, z) - - assert_allclose(v1, v2) diff --git a/tests/test_spatial.py b/tests/test_spatial.py index 94c4e6cb..580870b7 100644 --- a/tests/test_spatial.py +++ b/tests/test_spatial.py @@ -2,6 +2,7 @@ import pytest import numpy as np +import numpy.testing as npt from geoana import utils from geoana import spatial @@ -21,64 +22,157 @@ def test_errors(): class TestCoordinates(unittest.TestCase): - def test_rotate_vec_cyl2cart(self): - vec = np.r_[1., 0, 0].reshape(1, 3) - grid = np.r_[1., np.pi/4, 0].reshape(1, 3) - self.assertTrue(np.allclose( - spatial.cylindrical_to_cartesian(grid, vec), - np.sqrt(2)/2 * np.r_[1, 1, 0] - )) - self.assertTrue(np.allclose( - spatial.cylindrical_to_cartesian(grid), - np.sqrt(2)/2 * np.r_[1, 1, 0] - )) - self.assertTrue(np.allclose( - spatial.cartesian_to_cylindrical( - np.sqrt(2)/2 * np.r_[1, 1, 0] - ), - grid - )) - - self.assertTrue(np.allclose( - spatial.cartesian_to_cylindrical( - spatial.cylindrical_to_cartesian(grid) - ), - grid - )) - - self.assertTrue(np.allclose( - spatial.cartesian_to_cylindrical( - np.sqrt(2)/2 * np.r_[1, 1, 0], - spatial.cylindrical_to_cartesian(grid, vec) - ), - vec - )) - - vec = np.r_[0, 1, 2].reshape(1, 3) - grid = np.r_[1, np.pi/4, 0].reshape(1, 3) - self.assertTrue(np.allclose( - spatial.cylindrical_to_cartesian(grid, vec), - np.r_[-np.sqrt(2)/2, np.sqrt(2)/2, 2] - )) - - vec = np.r_[1., 0] - grid = np.r_[1., np.pi/4].reshape(1, 2) - self.assertTrue(np.allclose( - spatial.cylindrical_to_cartesian(grid, vec), - np.sqrt(2)/2 * np.r_[1, 1] - )) - def test_cartesian_to_cylindrical(self): - vec = np.r_[1., 0, 0] - grid = np.r_[1., np.pi / 4, 0] - grid_ = np.atleast_2d(grid) - vec_ = vec.reshape(grid_.shape, order='F') - theta = np.arctan2(grid_[:, 1], grid_[:, 0]) - c2c_test = np.hstack([utils.mkvc(np.cos(theta) * vec_[:, 0] + np.sin(theta) * vec_[:, 1], 2), - utils.mkvc(-np.sin(theta) * vec_[:, 0] + np.cos(theta) * vec_[:, 1], 2), - utils.mkvc(vec_[:, 2], 2)]) - c2c = spatial.cartesian_to_cylindrical(grid, vec) - np.testing.assert_equal(c2c_test, c2c) + cart_points = np.array([ + [1, 0, 0], + [0, 1, 0], + [0, 0, 1], + [-1, 0, 0], + [0, -1, 0], + [0, 0, -1], + [1, 1, 0], + ]) + + cyl_points = np.array([ + [1, 0, 0], + [1, np.pi/2, 0], + [0, 0, 1], + [1, np.pi, 0], + [1, -np.pi/2, 0], + [0, 0, -1], + [np.sqrt(2), np.pi / 4, 0] + ]) + + out_cyl = spatial.cartesian_to_cylindrical(cart_points) + np.testing.assert_allclose(out_cyl, cyl_points, atol=1E-15) + + out_cart = spatial.cylindrical_to_cartesian(cyl_points) + np.testing.assert_allclose(out_cart, cart_points, atol=1E-15) + + def test_cartesian_to_cylindrical_vector(self): + cart_points = np.array([ + [1, 0, 0], + [0, 1, 0], + [0, 0, 1], + [-1, 0, 0], + [0, -1, 0], + [0, 0, -1], + [1, 1, 0], + ]) + + cyl_vecs = np.array([ + [1, 0, 0], + [1, 0, 0], + [0, 0, 1], + [1, 0, 0], + [1, 0, 0], + [0, 0, -1], + [np.sqrt(2), 0, 0] + ]) + + cart_vecs2 = np.array([ + [0, 1, 0], + [-1, 0, 0], + [0, 0, 1], + [0, -1, 0], + [1, 0, 0], + [0, 0, -1], + [-1, 1, 0], + ]) + + cyl_vecs2 = np.array([ + [0, 1, 0], + [0, 1, 0], + [0, 0, 1], + [0, 1, 0], + [0, 1, 0], + [0, 0, -1], + [0, np.sqrt(2), 0], + ]) + + out_cyl = spatial.cartesian_to_cylindrical(cart_points, cart_points) + np.testing.assert_allclose(out_cyl, cyl_vecs, atol=1E-15) + + out_cyl = spatial.cartesian_to_cylindrical(cart_points, cart_vecs2) + np.testing.assert_allclose(out_cyl, cyl_vecs2, atol=1E-15) + + cyl_points = spatial.cartesian_to_cylindrical(cart_points) + + out_cart = spatial.cylindrical_to_cartesian(cyl_points, cyl_vecs) + np.testing.assert_allclose(out_cart, cart_points, atol=1E-15) + + out_cart = spatial.cylindrical_to_cartesian(cyl_points, cyl_vecs2) + np.testing.assert_allclose(out_cart, cart_vecs2, atol=1E-15) + + def test_cartesian_to_spherical(self): + test_points = np.array([ + [0, 0, 1], + [0, 1, 0], + [1, 0, 0], + [0, 0, -1], + [0, -1, 0], + [-1, 0, 0], + ]) + + sph_points = np.array([ + [1, 0, 0], + [1, np.pi/2, np.pi/2], + [1, 0, np.pi/2], + [1, 0, np.pi], + [1, -np.pi/2, np.pi/2], + [1, np.pi, np.pi/2], + ]) + + test_sph = spatial.cartesian_to_spherical(test_points) + + npt.assert_allclose(test_sph, sph_points, atol=1E-15) + + undo_test = spatial.spherical_2_cartesian(sph_points) + + npt.assert_allclose(undo_test, test_points, atol=1E-15) + + def test_cartesian_to_spherical_vector(self): + cart_points = np.array([ + [0, 0, 1], + [0, 1, 0], + [1, 0, 0], + [0, 0, -1], + [0, -1, 0], + [-1, 0, 0], + ]) + + cart_vec2 = np.array([ + [1, 0, 0], + [-1, 0, 0], + [0, -1, 0], + [-1, 0, 0], + [1, 0, 0], + [0, 1, 0], + ]) + + sphere_vecs2 = np.array([ + [0, 0, 1], + [0, 1, 0], # @ y=1, vec = -x_hat -> +phi_hat + [0, -1, 0], # @ x=1, vec = -y_hat -> -phi_hat + [0, 0, 1], + [0, 1, 0], # @ y=-1, vec=x-hat -> + phi_hat + [0, -1, 0], # @ x=-1, vec=y_hat -> - phi_hat + ]) + + # If we use these test points as the vectors to transform + # it should give us back vectors with only a radial component. + sphere_vecs = np.zeros_like(cart_points) + sphere_vecs[:, 0] = 1 + + test_vecs = spatial.cartesian_to_spherical(cart_points, cart_points) + npt.assert_allclose(test_vecs, sphere_vecs, atol=1E-15) + + test_vecs = spatial.cartesian_to_spherical(cart_points, cart_vec2) + npt.assert_allclose(test_vecs, sphere_vecs2, atol=1E-15) + + sphere_points = spatial.cartesian_to_spherical(cart_points) + out_cart = spatial.spherical_to_cartesian(sphere_points, sphere_vecs) + npt.assert_allclose(out_cart, cart_points, atol=1E-15) def test_vector_dot(self): xyz = utils.ndgrid(np.linspace(-1, 1, 20), np.array([0]), np.linspace(-1, 1, 20)) diff --git a/tests/test_utils.py b/tests/test_utils.py index 79b2a149..66b89b6e 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -20,12 +20,6 @@ def test_mkvc(): x_new = mkvc(x) np.testing.assert_equal(x_test, x_new) - y = np.matrix('1 2; 3 4') - y_test = np.array([[1, 2], [3, 4]]) - y_test = np.concatenate((y_test[:, 0], y_test[:, 1]), axis=None) - y_new = mkvc(y) - np.testing.assert_equal(y_test, y_new) - def test_nd_grid(): x = np.array([1])