How to Implement Analytical Magnetic Field Equations into a Physics-Informed Neural Network (PINN) in PyTorch?

Question

I'm working on a project that involves modeling the magnetic field around Europa, one of Jupiter's moons, using a Physics-Informed Neural Network (PINN) in PyTorch. My goal is to implement analytical expressions for the primary magnetic field and secondary induced field. At the moment, I'm trying to implement a simple version of the fields - to then more complex properties can be added to later on. The data used here are the Galileo Europa flybys, which contain magnetometer data in three axis along with the spacecraft coordinates.

Theory:

The primary field $B_{prim}$ from Jupiter can be obtained by a simple call to the Python library JupiterMag, which the model can then try to fit alongside the secondary field. It's mainly the secondary field we're focusing on implementing here.

For Europa, it is practical to consider it as a shell with uniform conductivity, $\sigma$, an inner radius $r_1$, and an outer radius $r_0 = r_1 + h$, which surrounds an insulating core and is itself encased by an insulating shell extending from $r_0$ to an outer radius $r_m = r_0 + d$ (where $r_m$ denotes the moon's radius).

We observe that within the conducting shell, the magnetic field's behavior is governed by a diffusion equation derived from Maxwell's equations and Ohm's law:

$\tag{1} \begin{equation} \nabla^2 \mathbf{B} = \mu \sigma \frac{\partial \mathbf{B}}{\partial t} \end{equation} $

Here, $\mu$ represents the permeability, which is assumed to be equal to that of a vacuum, $\mu_0$. Within the moon's insulating regions, Equation 1 reduces to Laplace's equation, i.e. the magnetic field satisfies:

$ \begin{equation} \tag{2} \nabla^2 \mathbf{B} = 0 \end{equation} $

In the surrounding space, where conductivity is negligible, Equation 2 also applies, disregarding plasma convection effects. The primary magnetic field, assumed to oscillate at a frequency $\omega$ in a direction defined by unit vector $\mathbf{e}_0$, is considered to be the real component of the complex vector:

$\tag{3} \begin{equation} \mathbf{B}_{prim} = B_{prim} e^{-i\omega t} \mathbf{e}_0 \end{equation} $

The boundary conditions for the total time-varying magnetic field $\mathbf{B}$ are as follows:

• $\mathbf{B}$ must be continuous across each shell's boundaries.
• At the core's center ($\mathbf{r} = 0$), $\mathbf{B}$ must not become infinite.
• As distance from the sphere increases ($|\mathbf{r}| \gg r_0$), $\mathbf{B}$ should asymptotically approach the external field, $\mathbf{B}_\text{prim}$.

We can express the total magnetic field as the sum of the primary and the secondary induced fields. For a simple layered shell model with uniform primary field and spherical symmetry in conductivity distribution, the analytical expression for $\mathbf{B}$ is obtainable. The secondary induced field outside the shell ($r > r_0$) can be modeled as a dipole field:

$\tag{4} \begin{equation} \mathbf{B}_{sec} = \frac{\mu_0}{4 \pi}\left[3(\mathbf{r} \cdot \mathbf{M}) \mathbf{r} - r^2 \mathbf{M}\right] / r^5, \end{equation} $

with the corresponding magnetic moment $\mathbf{M}$ sharing the oscillation frequency $\omega$ and direction $\mathbf{e}_0$ of the primary field:

$\tag{5} \begin{equation} \mathbf{M} = -\frac{4 \pi}{\mu_0} A e^{i \phi} \mathbf{B}_{prim} r_{m}^3 / 2 \end{equation} $

Substituting Equation 5 into Equation 4, we get the secondary field's formula which is an induced dipole with a dipole moment that is a phase lagged function of the primary field:

$\tag{6} \begin{equation} \mathbf{B}_{sec} = -A e^{-i(\omega t - \phi)} B_{prim} \left[3(\mathbf{r} \cdot \mathbf{e}_0) \mathbf{r} - r^2 \mathbf{e}_0\right] r_{m}^2 / (2 r^5) \end{equation} $

Where $A$ and $\phi$ can be represented by Bessel functions - but in our case, we're going to have the network predict their values.

Implementation:

I've been working on this using PyTorch, trying to implement three losses:

1. A physics-informed loss that enforces Equation (6), which represents the secondary magnetic field $\mathbf{B}_{\text{sec}}$ as influenced by Europa's conductivity.
2. A data-driven loss that minimizes the discrepancy between the network's total field $(B_{prim} + B_{sec})$ predictions and the observed magnetometer data from the Galileo flybys.
3. A continuity loss that ensures $\mathbf{B}$ is continuous across Europa's shell boundaries as specified by the boundary conditions.

The first two are quite straight forward, but the third one doesn't come easy to me.

My question is: has anyone here done something similar to this, or has input on how it's best to implement this in PyTorch or DeepXDE?

Here's an outline of what my model looks like at the moment, just to give an idea of the setup:

class Swish(nn.Module):
    def forward(self, x):
        return x * torch.sigmoid(x)

class MagneticFieldModel(nn.Module):
    def __init__(self):
        super(MagneticFieldModel, self).__init__()
        # Shared layers for feature extraction
        self.shared_layers = nn.Sequential(
            nn.Linear(3, 40),
            Swish(),
            nn.Linear(40, 40),
            Swish(),
        )
        self.B_prim_pred_output = nn.Linear(40, 3) # Output layer for the primary field prediction in xyz

        # Learnable parameters for the secondary field calculation
        self.A = nn.Parameter(torch.tensor([1.34]))  # Initial value for A
        self.phi = nn.Parameter(torch.tensor([0.0]))  # Initial value for phi

    def forward(self, x, t):
        shared_features = self.shared_layers(x)
        B_prim_pred = self.B_prim_pred_output(shared_features)

        # Compute the secondary field B_sec based on the formula
        B_sec = self.compute_B_sec(B_prim_pred, x, t)

        return B_prim_pred, B_sec

    def compute_B_sec(self, B_prim, coordinates, t):
        # Ensure coordinates, B_prim, and t are all float32 tensors
        omega = 0.00015
        coordinates = coordinates.to(torch.float32)
        B_prim = B_prim.to(torch.float32)
        t = t.to(torch.float32)
        
        # Compute r, the norm of the position vector
        r = torch.norm(coordinates, dim=1, keepdim=True)

        # Normalize r to get the unit direction vector e_0
        # get the direction of the primary field 
        e_0 = B_prim / torch.norm(B_prim, dim=0, keepdim=True)

        # Calculate the complex exponential term for each instance
        exp_term = torch.exp(-1j * (omega * t.unsqueeze(1) - self.phi))

        # Compute the vector term inside the brackets
        dot_product = torch.sum(e_0 * B_prim, dim=1, keepdim=True)
        term_in_brackets = 3 * dot_product * e_0 - (r**2) * e_0

        # Assemble B_sec according to the formula
        B_sec = -self.A * exp_term * B_prim * term_in_brackets * (r**2) / (2 * (r**5))

        # Taking the absolute value to deal with the complex number
        B_sec_magnitude = torch.abs(B_sec)

        # Returning the magnitude which should be a [32, 3] tensor
        return B_sec_magnitude.squeeze()


def compute_loss(model, coordinates, B_obs, B_prim_obs, t):
    # Forward pass to get A and B_prim_pred
    B_prim_pred, B_sec_pred = model(coordinates, t)

    B_total_pred = B_prim_pred + B_sec_pred

    # Compute each loss component
    B_prim_loss = F.mse_loss(B_prim_pred, B_prim_obs)
    B_total_loss = F.mse_loss(B_total_pred, B_obs)

    loss = B_prim_loss + B_total_loss
    #print(f"Loss components: B_prim: {B_prim_loss.item()}, B_total: {B_total_loss.item()}")
    return loss

Answer 1

Why not add a loss term of the form

$$ (\mathbf{B}^+ - \mathbf{B}^-)^2 $$

to your loss function? Here the superscripts $+,-$ denote two different sides of the shell interface.