mlff_attack.grad_based package

This package contains gradient-based adversarial attack implementations for MLFF models.

Submodules

mlff_attack.grad_based.fgsm module

Fast Gradient Sign Method (FGSM) attack implementation for MLFF models.

This module implements the FGSM attack specifically for MACE models, extending the base MLFFAttack class.

class mlff_attack.grad_based.fgsm.FGSM_MACE(model: Any, epsilon: float = 0.01, device: str = 'cpu', track_history: bool = True, target_energy: float | None = None)

Bases: MLFFAttack

FGSM attack implementation for MACE force field models.

The Fast Gradient Sign Method computes the gradient of the loss with respect to atomic positions and perturbs them in the direction that maximizes the loss.

model

MACE calculator attached to atoms

Type:

Any

epsilon

Step size for perturbation (Angstroms)

Type:

float

device

Device for PyTorch computations

Type:

str

target_energy

Optional target energy for the attack. If None, maximizes energy

Type:

float or None

__init__(model: Any, epsilon: float = 0.01, device: str = 'cpu', track_history: bool = True, target_energy: float | None = None)

Initialize FGSM , clip_radius: Optional[float] = Noneattack for MACE models.

Parameters:

• model (Any) – MACE calculator (will be attached to atoms)

• epsilon (float, optional) – Perturbation step size in Angstroms, by default 0.01

• device (str, optional) – Device for computations (‘cpu’ or ‘cuda’), by default ‘cpu’

• track_history (bool, optional) – Whether to track attack progression, by default True

• target_energy (Optional[float], optional) – Optional target energy (if None, maximize energy), by default None

attack(atoms: Any, n_steps: int = 1, clip: bool = True) → Any

Execute FGSM attack.

For standard FGSM, n_steps=1. For iterative FGSM (I-FGSM), use n_steps>1.

Parameters:

• atoms (Any) – Input atomic structure with MACE calculator

• n_steps (int, optional) – Number of attack iterations (1 for FGSM, >1 for I-FGSM), by default 1

• clip (bool, optional) – Whether to clip perturbations to epsilon bound, by default True

Returns:

Perturbed atoms object

Return type:

Any

attack_step(atoms: Any, step: int = 0) → Any

Perform one step of FGSM attack.

Parameters:

• atoms (Any) – Current atomic structure with MACE calculator

• step (int, optional) – Current iteration number, by default 0

Returns:

Updated atoms object with perturbed positions

Return type:

Any

compute_gradient(atoms: Any, loss_fn: Callable | None = None) → ndarray

Compute gradient of loss with respect to atomic positions.

Parameters:

• atoms (Any) – ASE Atoms object with MACE calculator

• loss_fn (Optional[Callable], optional) – Optional custom loss function. If None, uses default (maximize energy or target energy loss), by default None

Returns:

Gradient array with shape (n_atoms, 3)

Return type:

np.ndarray

get_attack_summary() → dict

Get a summary of the attack results.

Returns:

Dictionary with attack summary statistics

Return type:

dict

save_perturbation(filepath: str, atoms_original: Any | None = None, atoms_perturbed: Any | None = None, include_metadata: bool = True) → None

Save perturbation data with additional FGSM-specific information.

Parameters:

• filepath (str) – Output file path (.npz format)

• atoms_original (Optional[Any], optional) – Optional original atoms (for saving chemical symbols, cell), by default None

• atoms_perturbed (Optional[Any], optional) – Optional perturbed atoms, by default None

• include_metadata (bool, optional) – Whether to include attack parameters, by default True

mlff_attack.grad_based.mlff_attack_class module

class mlff_attack.grad_based.mlff_attack_class.MLFFAttack(model: Any, epsilon: float, device: str = 'cpu', track_history: bool = True)

Bases: ABC

Base class for adversarial attacks on Machine Learning Force Fields.

This abstract class defines the interface for implementing various attack strategies on MLFF models like MACE, ALIGNN, etc.

model

The MLFF model (e.g., MACE calculator, ALIGNN ForceField)

Type:

Any

epsilon

Maximum perturbation magnitude (in Angstroms for position attacks)

Type:

float

device

Device for computations (‘cpu’, ‘cuda’, etc.)

Type:

str

attack_history

Dictionary storing attack trajectory information including energies, forces, perturbations, and gradients

Type:

dict or None

__init__(model: Any, epsilon: float, device: str = 'cpu', track_history: bool = True)

Initialize the attack.

Parameters:

• model (Any) – MLFF model with calculator interface

• epsilon (float) – Maximum perturbation magnitude

• device (str, optional) – Device for PyTorch computations, by default ‘cpu’

• track_history (bool, optional) – Whether to track attack progression, by default True

attack(atoms: Any, n_steps: int = 1, clip: bool = True) → Any

Execute the complete attack.

Parameters:

• atoms (Any) – Input atomic structure to attack

• n_steps (int, optional) – Number of attack iterations, by default 1

• clip (bool, optional) – Whether to clip perturbations to epsilon bound, by default True

Returns:

Perturbed atoms object

Return type:

Any

abstractmethod attack_step(atoms: Any, step: int = 0) → Any

Perform one step of the adversarial attack.

Parameters:

• atoms (Any) – Current atomic structure

• step (int, optional) – Current iteration number, by default 0

Returns:

Updated atoms object with perturbed positions

Return type:

Any

abstractmethod compute_gradient(atoms: Any, loss_fn: Callable | None = None) → ndarray

Compute gradient of loss with respect to atomic positions.

Parameters:

• atoms (Any) – ASE Atoms object or equivalent structure

• loss_fn (Optional[Callable], optional) – Optional custom loss function (default: maximize energy), by default None

Returns:

Gradient array with shape (n_atoms, 3)

Return type:

np.ndarray

get_perturbation_stats() → Dict[str, float]

Get statistics about the current perturbation.

Returns:

Dictionary with perturbation statistics

Return type:

Dict[str, float]

load_perturbation(filepath: str) → Dict[str, ndarray]

Load perturbation data from file.

Parameters:

filepath (str) – Input file path (.npz format)

Returns:

Dictionary containing loaded data

Return type:

Dict[str, np.ndarray]

reset() → None

Reset attack state.

save_perturbation(filepath: str, include_metadata: bool = True) → None

Save perturbation data to file.

Parameters:

• filepath (str) – Output file path (.npz format)

• include_metadata (bool, optional) – Whether to include attack parameters, by default True

mlff_attack.grad_based.pgd module

Projected Gradient Descent (PGD) attack implementation for MLFF models.

This module implements the PGD attack specifically for MACE models, extending the base MLFFAttack class.

class mlff_attack.grad_based.pgd.PGD_MACE(model: Any, epsilon: float, alpha: float, num_iter: int, device: str = 'cpu', track_history: bool = True)

Bases: MLFFAttack

__init__(model: Any, epsilon: float, alpha: float, num_iter: int, device: str = 'cpu', track_history: bool = True)

Initialize the PGD attack.

Parameters:

• model – MLFF model with calculator interface

• epsilon – Maximum perturbation magnitude

• alpha – Step size for each iteration

• num_iter – Number of attack iterations

• device – Device for PyTorch computations

• track_history – Whether to track attack progression

attack() → Any

Execute the full PGD attack over the specified number of iterations.

Returns:

Final perturbed atomic structure after all attack iterations

attack_step(atoms: Any, step: int = 0) → Any

Perform one step of the PGD adversarial attack.

Parameters:

• atoms – Current atomic structure

• step – Current iteration number

Returns:

Updated atomic structure after one attack step

compute_gradient(atoms: Any, loss_fn: Callable | None = None) → ndarray

Compute gradient of loss with respect to atomic positions.

Parameters:

• atoms – ASE Atoms object or equivalent structure

• loss_fn – Optional custom loss function (default: maximize energy)

Returns:

Gradient array with shape (n_atoms, 3)

Module contents