| |
| |
|
|
| from functools import partial |
| from typing import Dict, Literal, Optional |
|
|
| from ..diffusion.losses import SummedFieldLoss, denoising_score_matching |
| from ..diffusion.model_target import ModelTarget |
| from ..diffusion.training.field_loss import FieldLoss, d3pm_loss |
| from ..diffusion.wrapped.wrapped_normal_loss import wrapped_normal_loss |
|
|
|
|
| class MaterialsLoss(SummedFieldLoss): |
| def __init__( |
| self, |
| reduce: Literal["sum", "mean"] = "mean", |
| d3pm_hybrid_lambda: float = 0.0, |
| include_pos: bool = True, |
| include_cell: bool = True, |
| include_atomic_numbers: bool = True, |
| weights: Optional[Dict[str, float]] = None, |
| ): |
| model_targets = {"pos": ModelTarget.score_times_std, "cell": ModelTarget.score_times_std} |
| self.fields_to_score = [] |
| self.categorical_fields = [] |
| loss_fns: Dict[str, FieldLoss] = {} |
| if include_pos: |
| self.fields_to_score.append("pos") |
| loss_fns["pos"] = partial( |
| wrapped_normal_loss, |
| reduce=reduce, |
| model_target=model_targets["pos"], |
| ) |
| if include_cell: |
| self.fields_to_score.append("cell") |
| loss_fns["cell"] = partial( |
| denoising_score_matching, |
| reduce=reduce, |
| model_target=model_targets["cell"], |
| ) |
| if include_atomic_numbers: |
| model_targets["atomic_numbers"] = ModelTarget.logits |
| self.fields_to_score.append("atomic_numbers") |
| self.categorical_fields.append("atomic_numbers") |
| loss_fns["atomic_numbers"] = partial( |
| d3pm_loss, |
| reduce=reduce, |
| d3pm_hybrid_lambda=d3pm_hybrid_lambda, |
| ) |
| self.reduce = reduce |
| self.d3pm_hybrid_lambda = d3pm_hybrid_lambda |
| super().__init__( |
| loss_fns=loss_fns, |
| weights=weights, |
| model_targets=model_targets, |
| ) |
|
|
