Indexed codewriting #22219
ThePauliPrinciple
started this conversation in
Show and tell
Indexed codewriting
#22219
Replies: 1 comment 1 reply
-
In general this sounds good. I think that maybe there should be a different version of |
Beta Was this translation helpful? Give feedback.
1 reply
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
-
For a few years already, I have been interrested in using (and have used) Sympy's code writing for indexed objects. In particular, I have used machine learning for obtaining continuous representations of (tensorial) properties of atoms, such as energies, forces and friction tensors, needed for simulations of the dynamics of atoms during chemical reactions.
Since the properties of such systems are bound by symmetries (permutations of the same atoms as well as total rotational symmetry and translation symmetry for periodic systems), the neural networks do not use the cartesian coordinates as inputs directly, but instead use "symmetry functions" which ensure the same input the NNs see is the same if symmetry equivalent cartesian coordinates are given for the atoms in the system. One such choice can be found in equations S1-S4 in this sup. inf. : https://pubs.acs.org/doi/abs/10.1021/acs.jpclett.7b00784. Symmetry functions are a bit of a "flavour of the month" choice, so it is convenient if code can be automatically generated using Sympy, which is already possible, although not without modifications.
Moreover, for tensorial properties of these atoms, jacobians or hessians are required to obtain the correct symmetry of the output (e.g. if two atoms swap, the forces should swap as in, it should be permutation equivariant not invariant). Obtaining jacobians and hessions from these symmetry functions is a lot easier with Sympy than by hand.
For convenience, I want to be able to convert the expressions to tensorflow code for fitting, while c or fortran code is required for integration into the academic spaghetti stack. Currently, Indexed objects are not supported for tensorflow code writing at all, and for good reason. It is not trivial at all to correctly write an Indexed expression into a tensorflow expression, as there might be a lot axis swapping or new axis creation (for broadcasting) required. Fortunately, I have made an attempt to create this functionality.
In particular, I have written the class "ExplicitIndexedAssignment" inheriting from the Assignment class from sympy.codegen.ast.
It is currently defined as "equate the lhs to the sum over all indices not in the lhs of the rhs". In particular, indexed objects are not considered tensors in this assignment, that is to say, if a multiplication contains two the same indices, it is NOT considered a contraction. Currently it works as follows:
Note here in particular that it detected it needs to create a new axis for the subtraction to make sense in tensorflow:
The expression yields a matrix of atom-atom distances, as are often used in symmetry functions.
My question here is if people feel this could be a good addition to sympy, or if I should write this in my own library, since it does not add anything to the mathematics in sympy, only to the code writing.
There were some modifications I had to make to sympy:
There are several features I would still like to implement:
All in all, I'm looking to write sympy expressions which with the Latex printer result in similar expressions as can be found in research papers, but then can be immediately written as tensorflow, c or fortran code, even if that code might not be the most optimized code.
Feel free to share any ideas or suggestions.
Beta Was this translation helpful? Give feedback.
All reactions