File size: 2,065 Bytes
f15d29e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 | # Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.
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,
)
|