-
Notifications
You must be signed in to change notification settings - Fork 4
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Tensornet representation #167
Conversation
…y, the array of atomic species was assumed to be unique; now we consider the order of the species array as well. Only entries with the same atomic species array that appear in order are considered to be the same molecule for purposes of grouping conformers. That is, if the same species array is encountered later in the array, it is treated as a new molecule. Molecule names are not just the string made from their species array, but also have a molecule number appended to them.
This revises how unique molecules are determined for ANI2x and updates dataset yaml files.
@wiederm
|
@wiederm Just as a reminder, |
@wiederm |
@wiederm The only change I made that may affects 'main' branch is that 'CosineCutoff' now has a 'representation_unit' parameter that is set default with 'unit.nanometer'. (So I would assume that this won't even change anything outside tensornet.)
|
We discussed this offline: edge_weight is the pairwise distance, and edge_vec is the pairwise distance vector. Box vector and batch are passed to the input and don't need to be generated. The data structure is the
So does modelforge! All pairwise interactions are listed, that means every atom |
No, ANI uses angstrom. Testing is pretty simple, if you have a class that implements some transformation you can simply do: input_data_in_angstrom = torch.abs(torch.randn(5,1)) * 5 # generate random data on interval [0,5]
reference_implementation = RefClass()
modelforge_implementation = MfClass()
# the rbf is defined on an interval [min_distance, max_distance].
# you can convert between rbfs in different length units by scaling the rbfs,
# therefore you can convert from nm to A by multiplying with 10
assert torch.allclose(modelforge_implementation(input_data_in_angstrom/10)*10, reference_implementation(input_data_in_angstrom)) |
@wiederm In class 'ExpNormalSmearing':
The output of this forward function is called 'edge_attr' in 'TensorEmbedding'. This should be However, in 'TensorEmbedding':
Cutoff function C is timed again to I, A, and S. |
@wiederm
|
…rent output from nn.Linear)
@wiederm In modelforge, if we have 3 atoms The real problem here is that in torchmd-net, since the first 6 pairs are the same as the last 6 pairs, I assume the output of that linear layer should also be the case (the first 6 tensors are the same as the last 6 tensors), while it's not the case. Thus in later summation, results from two ways are different. One way that can definitely fix this problem is to convert the modelforge index system to the torchmd-net one inside representation module. But I think that may not be what we want in modelforge. We should discuss how we can solve the problem. |
Note the boolean set in the Pairlist: modelforge/modelforge/potential/models.py Line 56 in 5a312b0
If this is set to |
I am not compley sure that I follow this:
Can you provide a minimum example that examplifies your concern? |
If possible, it might be very useful to discuss this in a working test |
…ensornet_representation # Conflicts: # modelforge/potential/tensornet.py # modelforge/tests/test_tensornet.py
Thanks! It has been fixed by now. |
@@ -10,11 +14,23 @@ | |||
from modelforge.potential.utils import TensorNetRadialSymmetryFunction | |||
from modelforge.potential.utils import NeuralNetworkData | |||
from modelforge.potential.utils import NNPInput | |||
from .models import PairListOutputs | |||
|
|||
ATOMIC_NUMBER_TO_INDEX_MAP = { |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I'd suggest importing this from potential/ani2x.py
self.rsf_projection_I = nn.Linear( | ||
number_of_radial_basis_functions, | ||
hidden_channels, | ||
dtype=dtype, |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
you don't have to set dtype
here. All parameters in sub-models are automatically registered as nn.parameters
and device and type can be set for each of the parameter at runtime.
self.rsf_projection_A = nn.Linear( | ||
number_of_radial_basis_functions, | ||
hidden_channels, | ||
dtype=dtype, | ||
) | ||
self.rsf_projection_S = nn.Linear( | ||
number_of_radial_basis_functions, | ||
hidden_channels, | ||
dtype=dtype, | ||
) | ||
self.atomic_number_i_embedding_layer = nn.Embedding( | ||
max_atomic_number, | ||
hidden_channels, | ||
dtype=dtype, | ||
) | ||
self.atomic_number_ij_embedding_layer = nn.Linear( | ||
2 * hidden_channels, | ||
hidden_channels, | ||
dtype=dtype, | ||
) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
it might make sense to use a ModuleDict
here: https://pytorch.org/docs/stable/generated/torch.nn.ModuleDict.html
self.rsf_projections = ModuleDict(
{ A : n.Linear(
number_of_radial_basis_functions,
hidden_channels,
dtype=dtype,
),
B: ...,
)
for _ in range(3): | ||
self.linears_tensor.append( | ||
nn.Linear(hidden_channels, hidden_channels, bias=False) | ||
) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
are you sure that for each lineare layer bias=False
?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
and no activation function is used?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I'd prefer not to use append
operations for `ModuleList. It makes it difficult to identify all to components of the list at one glance. In that case I'd suggest to use something like:
self.linears_tensor = nn.ModuleList(
[nn.Linear(hidden_channels, hidden_channels, bias=False) for _ in range(3)]
)
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
maybe use [nn.Sequential]
(https://pytorch.org/docs/stable/generated/torch.nn.Sequential.html)?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
and maybe the Dense implementation of modelforge?
self.linears_scalar.append( | ||
nn.Linear(hidden_channels, 2 * hidden_channels, bias=True, dtype=dtype) | ||
) | ||
self.linears_scalar.append( | ||
nn.Linear(2 * hidden_channels, 3 * hidden_channels, bias=True, dtype=dtype) | ||
) | ||
self.init_norm = nn.LayerNorm(hidden_channels, dtype=dtype) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
You can write this as:
self.linear_layers_for_scalar_properties = nn.ModuleList(
nn.Linear(hidden_channels, 2 * hidden_channels, bias=True),
nn.Linear(2 * hidden_channels, 3 * hidden_channels, bias=True)
)
which is slightly easier to read.
GaussianRadialBasisFunction(), | ||
representation_unit: unit.Quantity = unit.angstrom, | ||
): | ||
self.representation_unit = representation_unit |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I don't think that this is necessary. All internal units are in nanometer, but for certain operations it might be useful to convert to angstrom.
_max_distance = self.max_distance.to(self.representation_unit).m | ||
_min_distance = self.min_distance.to(self.representation_unit).m |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
convert everything to nanometer
): | ||
# TensorNet evenly distribute centers in log space | ||
start_value = torch.exp( | ||
torch.scalar_tensor( | ||
-_max_distance_in_nanometer*10 + _min_distance_in_nanometer*10, # in angstrom | ||
-_max_distance + _min_distance, # in angstrom |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
You can write this exactly as outlined here:
modelforge/modelforge/potential/utils.py
Line 810 in 5b0e9bc
def calculate_radial_basis_centers( |
@@ -822,31 +857,35 @@ def calculate_radial_basis_centers( | |||
return centers | |||
|
|||
def calculate_radial_scale_factor( |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
at first glance this should also match the PhysNetRadialBasisFunction. But please double check that claim!
_max_distance = self.max_distance.to(self.representation_unit).m | ||
_min_distance = self.min_distance.to(self.representation_unit).m | ||
alpha = 5.0 / (_max_distance - _min_distance) | ||
d_ij = torch.exp(alpha * (-d_ij + _min_distance)) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
d_ij
should always be the pairwise distance. If transformations are performed that are specific for the radial basis functions I think these should be moved inside the radial basis function implementation (and shouldn't overwrite d_ij
).
Was this resolved? It seems that they wanted to make sure that the outputs of intermediate linear layers also decay to 0 at the cutoff distance, even before they multiply by the cutoff-filtered distances. |
Yes, this has been resolved! We came to the same conclusion and left it in the code. |
Closing PR since it is superseeded by PR #181 |
Description
Implement modelforge TensorNetRepresentation module, which matches the output of TensorEmbedding forward. In original implementation, TensorEmbedding is used by TensorNet class, where forward is called whenever updating TensorNet. The output of TensorEmbedding corresponds to "X" tensor in the TensorNet paper.
Todos
Questions
Status