Magnetic potential/field improvements for collinear GPAW calculations

[TomaSusi]

Summary

• Fixed FieldArray.tile() for vector-valued fields (VectorPotentialArray, MagneticFieldArray): it hard-coded the last 3 array axes as (z, x, y), which only holds for scalar potentials, so tiling a vector/magnetic field silently tiled the wrong axis and raised a slice-thickness mismatch error.
• GPAWMagneticField/GPAWVectorPotential now raise NotImplementedError for frozen_phonons/repetitions instead of silently ignoring them (they were accepted but never implemented).
• rotate_field now defaults to "auto": since collinear spin has no real-space direction, GPAW's internal magnetization axis always produces Az == 0. "auto" automatically swaps whichever in-plane component (Ax or Ay) has the larger magnitude into z, restricted to the two physically meaningful 90-degree swaps (not an arbitrary blended angle). Explicit rotate_field tuples and rotate_field=None (raw output) still work as before.
• Added gpaw_magnetic_fields(...), which builds the electrostatic potential, vector potential, and (optionally) magnetic field from the same GPAW calculator(s) in one call, returning a GPAWMagneticFields bundle with .tile() and .combined_potential() to support combining a rigid ab initio magnetic contribution with a separately tiled/frozen-phonon-ensembled electrostatic potential.
• Added GPAWMagneticFields.show() for a side-by-side V/Ax/Az/Bx/Bz comparison plot.

Usage

import numpy as np
import gpaw

import abtem
from abtem.inelastic.phonons import FrozenPhonons
from abtem.magnetism.gpaw import gpaw_magnetic_fields

# ------------------------------------------------------------------
# 1. Load the converged, spin-polarized (collinear) DFT calculation.
# ------------------------------------------------------------------
atoms, calc = gpaw.restart("saved_calculation.gpw")

sampling = 0.05
energy = 300e3          # electron energy [eV]
reps = (5, 5, 4)         # tile the DFT unit cell into a crystal

# ------------------------------------------------------------------
# 2. Build the electrostatic potential, vector potential, and (if
#    needed) magnetic field from the same calculator in one call.
#
#    rotate_field defaults to "auto": it picks whichever in-plane
#    component (Ax or Ay) is larger and swaps it into Az, so Az is
#    guaranteed nonzero without any manual Euler-angle bookkeeping.
#    Pass rotate_field=(0.0, np.pi / 2, 0.0) (or any other tuple) to
#    force a specific orientation instead, or rotate_field=None to see
#    the raw (Az == 0) output.
# ------------------------------------------------------------------
fields = gpaw_magnetic_fields(
    calc,
    sampling=sampling,
    plane="xy",
    include_magnetic_field=True,   # set False (default) to skip B; only
                                    # needed here for visualization below
)

fields.potential          # PotentialArray  (electrostatic, exact DFT config)
fields.vector_potential   # VectorPotentialArray, shape (num_slices, 3, *gpts)
fields.magnetic_field     # MagneticFieldArray, or None if not requested

fields.show()                    # V, Ax, Az, Bx, Bz side by side
fields.show(tile=(2, 2))         # preview a tiled/repeated view

# ------------------------------------------------------------------
# 3a. Simple case: no frozen phonons, single deterministic potential.
#     Combine at the unit-cell level, then tile the combined result.
# ------------------------------------------------------------------
magnetic_potential_unit = fields.combined_potential(energy=energy)

magnetic_crystal_potential = (
    abtem.CrystalPotential(magnetic_potential_unit, repetitions=reps)
    .build()
    .compute()
)

# ------------------------------------------------------------------
# 3b. Frozen-phonon case: use rattled IAM configurations for the
#     electrostatic part (thermal diffuse scattering), while keeping
#     the ab initio magnetic vector potential rigid across all of
#     them. Combination now has to happen *after* tiling, since the
#     electrostatic ensemble and the magnetic part are tiled
#     independently and only line up once both are the crystal size.
# ------------------------------------------------------------------
frozen_phonons = FrozenPhonons(atoms, num_configs=8, sigmas=0.1)
electrostatic_unit = abtem.Potential(
    frozen_phonons, sampling=sampling, slice_thickness=1.0
)

# num_frozen_phonons independent tiled realizations; each tile of each
# realization draws from the 8-configuration displacement pool.
crystal_potential = (
    abtem.CrystalPotential(
        electrostatic_unit,
        repetitions=reps,
        num_frozen_phonons=6,
        ensemble_mean=False,
    )
    .build()
    .compute()
)

# Tile the (rigid, non-ensemble) magnetic components to the same size.
tiled_fields = fields.tile(reps)

# Now combine: adjust_coulomb_potential broadcasts the single tiled
# vector potential against every frozen-phonon realization.
magnetic_crystal_potential = tiled_fields.combined_potential(
    energy=energy, potential=crystal_potential
)
# -> PotentialArray with ensemble_shape == (6,), ready for multislice.

# ------------------------------------------------------------------
# 4. Use exactly like any other potential downstream -- no special
#    casing needed since combined_potential() already returns a
#    normal PotentialArray.
# ------------------------------------------------------------------
# waves = abtem.PlaneWave(energy=energy, sampling=sampling)
# exit_waves = waves.multislice(magnetic_crystal_potential).compute()

Test plan

• Verified against real spin-polarized GPAW calculations (GPAW 26.7.1b1) that tile(), the auto rotation, gpaw_magnetic_fields(), and combined_potential() are numerically correct/bit-identical to the equivalent manual API calls.
• Ran the exact usage example above (both the simple and frozen-phonon paths) end to end against a real GPAW 26.7.1b1 calculation.
• Verified GPAWMagneticFields.show() renders correctly (including the fix for it doubly-rendering in Jupyter).
• test/test_magnetics.py, test/test_potentials.py pass (42 passed, 18 skipped).

🤖 Generated with Claude Code