Python API

• nequip.data
  • Data Fields
    • _GRAPH_FIELDS
    • _NODE_FIELDS
    • _EDGE_FIELDS
    • _LONG_FIELDS
    • _CARTESIAN_TENSOR_FIELDS
  • nequip.data.datamodule
    • NequIPDataModule
      • NequIPDataModule.get_statistics()
    • ASEDataModule
    • sGDML_CCSD_DataModule
    • rMD17DataModule
    • MD22DataModule
    • NequIP3BPADataModule
    • COLLDataModule
    • TM23DataModule
    • SAMD23DataModule
    • WaterDataModule
  • nequip.data.dataset
    • AtomicDataset
      • AtomicDataset.num_atoms()
    • NequIPLMDBDataset
      • NequIPLMDBDataset.save_from_iterator()
    • ASEDataset
    • HDF5Dataset
    • NPZDataset
    • EMTTestDataset
    • SubsetByRandomSlice
    • RandomSplitAndIndexDataset()
  • nequip.data.transforms
    • ChemicalSpeciesToAtomTypeMapper
      • ChemicalSpeciesToAtomTypeMapper.forward()
    • NeighborListTransform
      • NeighborListTransform.forward()
    • VirialToStressTransform
      • VirialToStressTransform.forward()
    • StressSignFlipTransform
      • StressSignFlipTransform.forward()
    • AddNaNStressTransform
      • AddNaNStressTransform.forward()
  • Data Modifiers
    • PerAtomModifier
    • MappedFieldModifier
    • EdgeLengths
    • NumNeighbors
  • Dataset Statistics
    • CommonDataStatisticsManager()
    • EnergyOnlyDataStatisticsManager()
    • DataStatisticsManager
      • DataStatisticsManager.reset()
      • DataStatisticsManager.get_statistics()
    • Mean
    • MeanAbsolute
    • RootMeanSquare
    • StandardDeviation
    • Min
    • Max
    • Count
• nequip.train
  • NequIPLightningModule
    • NequIPLightningModule
    • EMALightningModule
    • ConFIGLightningModule
    • EMAConFIGLightningModule
    • ScheduleFreeLightningModule
  • Loss Function and Error Metrics
    • Simplified Wrappers
      • EnergyForceLoss
      • EnergyForceStressLoss
      • EnergyForceMetrics
      • EnergyForceStressMetrics
      • EnergyOnlyLoss
      • EnergyOnlyMetrics
    • Advanced Configuration: MetricsManager
      • MetricsManager
        • MetricsManager.parse_entry()
        • MetricsManager.forward()
        • MetricsManager.compute()
        • MetricsManager.set_coeffs()
    • Error Metrics
      • MeanSquaredError
      • RootMeanSquaredError
      • MeanAbsoluteError
      • MaximumAbsoluteError
      • HuberLoss
      • StratifiedHuberForceLoss
  • nequip.train.callbacks
    • SoftAdapt
    • LossCoefficientScheduler
    • LinearLossCoefficientScheduler
    • LossCoefficientMonitor
    • TestTimeXYZFileWriter
    • ValTimeXYZFileWriter
    • WandbWatch
    • TF32Scheduler
  • nequip.train.SimpleDDPStrategy
    • SimpleDDPStrategy
      • SimpleDDPStrategy.post_backward()
  • Custom Optimizers
    • MuonWithAuxAdam
      • MuonWithAuxAdam
• nequip.nn
  • ZBL
  • PerTypeScaleShift
    • PerTypeScaleShift.modify_PerTypeScaleShift()
• nequip.model
  • NequIP Message Passing GNN Models
    • NequIPGNNModel()
    • Preset NequIP Message Passing GNN Model
      • PresetNequIPGNNModel()
  • Saved Models
    • Loading Saved Models
      • ModelFromCheckpoint()
      • ModelFromPackage()
    • Modifying Saved Models
      • modify()
  • Parameter Groups for Optimizers
    • MuonParamGroups()
• nequip.integrations.ase
  • NequIPCalculator
    • NequIPCalculator.implemented_properties
  • NequIPCalculator.from_compiled_model()
• nequip.integrations.torchsim
  • NequIPTorchSimCalc
    • NequIPTorchSimCalc.compute_forces
    • NequIPTorchSimCalc.compute_stress
    • NequIPTorchSimCalc.from_compiled_model()
    • NequIPTorchSimCalc.setup_from_system_idx()
    • NequIPTorchSimCalc.forward()