Hi!

ForceNet was originally developed outside the OCP repo but the model was later ported over upon submission. It was benchmarked in the current repo to ensure results are consistent, but like you mentioned may be missing some things.

Regarding the transformation, the original implementation is shared below - `Random3DRotate`

:

```
class Random3DRotate(object):
def __init__(self, axis=None):
r"""
Rotate both pos and force
args:
axis (int): select which axis to rotate along. default: None (randomly rotate along all the axis)
"""
self.axis = axis
def __call__(self, data):
if self.axis is None:
# rotate along x-axis
degree = math.pi * random.uniform(-180, 180) / 180.0
sin, cos = math.sin(degree), math.cos(degree)
matrix = [[1, 0, 0], [0, cos, sin], [0, -sin, cos]]
data.pos = torch.matmul(
data.pos, torch.tensor(matrix).to(data.pos.dtype).to(data.pos.device)
)
data.force = torch.matmul(
data.force,
torch.tensor(matrix).to(data.force.dtype).to(data.force.device),
)
data.cell = torch.matmul(
data.cell,
torch.tensor(matrix).to(data.force.dtype).to(data.force.device),
)
# rotate along y-axis
degree = math.pi * random.uniform(-180, 180) / 180.0
sin, cos = math.sin(degree), math.cos(degree)
matrix = [[cos, 0, -sin], [0, 1, 0], [sin, 0, cos]]
data.pos = torch.matmul(
data.pos, torch.tensor(matrix).to(data.pos.dtype).to(data.pos.device)
)
data.force = torch.matmul(
data.force,
torch.tensor(matrix).to(data.force.dtype).to(data.force.device),
)
data.cell = torch.matmul(
data.cell,
torch.tensor(matrix).to(data.force.dtype).to(data.force.device),
)
# rotate along z-axis
degree = math.pi * random.uniform(-180, 180) / 180.0
sin, cos = math.sin(degree), math.cos(degree)
matrix = [[cos, sin, 0], [-sin, cos, 0], [0, 0, 1]]
data.pos = torch.matmul(
data.pos, torch.tensor(matrix).to(data.pos.dtype).to(data.pos.device)
)
data.force = torch.matmul(
data.force,
torch.tensor(matrix).to(data.force.dtype).to(data.force.device),
)
data.cell = torch.matmul(
data.cell,
torch.tensor(matrix).to(data.force.dtype).to(data.force.device),
)
else:
axis = self.axis
degree = math.pi * random.uniform(-180, 180) / 180.0
sin, cos = math.sin(degree), math.cos(degree)
if axis == 0:
matrix = [[1, 0, 0], [0, cos, sin], [0, -sin, cos]]
elif axis == 1:
matrix = [[cos, 0, -sin], [0, 1, 0], [sin, 0, cos]]
else:
matrix = [[cos, sin, 0], [-sin, cos, 0], [0, 0, 1]]
data.pos = torch.matmul(
data.pos, torch.tensor(matrix).to(data.pos.dtype).to(data.pos.device)
)
data.force = torch.matmul(
data.force,
torch.tensor(matrix).to(data.force.dtype).to(data.force.device),
)
data.cell = torch.matmul(
data.cell,
torch.tensor(matrix).to(data.force.dtype).to(data.force.device),
)
return data
```

The transformation happens at the dataloader level, which by default sets `transform=None`

. To modify this, you need to define the `transform`

here to:

```
self.train_dataset = registry.get_dataset_class(
self.config["task"]["dataset"]
)(self.config["dataset"], Random3DRotate())
```

FWIW - The repo currently has a similar transformation. This is missing the force rotation as part of the call for it to be useable for training, but I can push that change soon.

Hope this helps! Let me know if you run into any issues.