Torsion wrapper

[joshkamm]

I'm working on an API that wraps torsion functionality onto stk analogous to what is currently implemented for atoms and bonds.

It's under construction but looks roughly like this:

class MoleculeTorsioned(stk.Molecule):
  def get_torsions():
    'yield the torsions in the molecule'
  def get_torsion_list():
    'return a list of torsions in the molecule, where each torsion is a list of atom indices'

class ConstructedMoleculeTorsioned(stk.ConstructedMolecule):
  def get_torsion_infos():
    'yield data about the torsions in the molecule'
  def get_torsions_by_building_block():
    'return a set of the torsions corresponding to each building block of this molecule'

class TorsionInfo:
  def get_torsion():
    'get the torsion about which information is held.'
  def get_building_block():
    'get the building block from which the torsion originates'
  def get_building_block_id():
    'get the id of the torsion's building block'
  def get_building_block_torsion():
    'return the corresponding torsion in the building block where the torsion came from'

class Torsion(atom1, atom2, atom3, atom4)
  'represents a torsion within a molecule'

I'm working on these classes as wrappers but @lukasturcani, if any of this seems like it would fall within the scope of the stk project I could try contributing that to the stk repo itself rather than as an external wrapper. Otherwise I thought I'd post here in case anyone else finds it helpful.

[lukasturcani]

Hi Josh! This is awesome. Yeah, I'm kind of torn. I like the idea of having stk quite focused on specific functionality, notably constructing stuff, depositing that stuff into databases, and automatically constructing stuff, rather than providing tools for molecular analysis. My reasoning for this is that there is a lot of analysis that someone can do, so picking which functionality to add and which not to can be hard to do consistently. If this is something that can be added by using only the public methods currently available, then I think this works best as a separate library. Notably https://github.com/JelfsMaterialsGroup/stko takes this approach.

I guess one of my reasons is that there is the possibility of treading into territory of libraries, like rdkit, which already do a fair amount of analysis, so I'd rather focus stk development on things that only it can do. Also, having stk extensions like this managed outside the library allows their authors more control over the direction and content of their work.

That being said, I am happy to provide a link in the ReadMe and stk docs, and indeed usage examples in the docs, to any projects that use stk and are useful to users of stk, which I believe this would be. (Or accept PRs which add this stuff to the docs / readme)

Also happy to provide input for the API / implementation, and if stk doesn't allow you to build this API then happy to discuss how that can be achived, because I think is exactly the sort of thing that stk should allow you to build (and build simply / straightforwardly).

[joshkamm]

Thanks! Based on your comment, Andrew's comment, and looking around it seems like you're focusing stk on core molecular construction and that your research group has an stk extension (stko) for optimization / analysis which makes sense. It seems like contributing these wrappers to stko may be helpful.

One of the unknowns I had in mind was how difficult it would be to subclass the stk molecular classes effectively, particularly when there are multiple layers of construction (when a constructed molecule is used as a building block for further construction). However, if a goal of stk is to be easily extensible then implementing these extensions externally seems like it could actually be a good test of that. I can let you know if I run into questions or issues with how stk molecules behave when subclassed.

[lukasturcani]

Normally I would expect the multiple layers of construction to be done by an explicit conversion of a ConsturctedMolecule into a BuildingBlock with something like

constructed_mol = stk.ConstructedMolecule(...)
constructed_mol_bb = stk.BuildingBlock.init_from_molecule(constructed_mol)

So I would normally expect a ConsturctedMolecule to be "reduced" back down to BuildingBlock before doing further construction with it. This should, ideally, keep everything flat and simple. I wouldn't expect subclassing to be used here. If I was writing user code with hierarchical construction I might structure it something like (though this is not necessarily the right or best way):

def get_bb_alpha(bb1: stk.BuildingBlock, bb2: stk.BuildingBlock) -> stk.BuildingBlock:
    contructed_bb = stk.ConsructedMolecule(
        topology_graph=stk.SomeTopologyGraph(
            building_blocks=(bb1, bb2),
       ),
    )
    return stk.BuildingBlock.init_from_molecule(constructed_bb, [stk.SomeFactory()])

def get_complex_molecule(bb1: stk.BuildingBlock, bb2: stk.BuildingBlock) -> stk.ConstructedMolecule:
    return stk.ConstructedMolecule(
        topology_graph=stk.OtherTopologyGraph(
            bb1=get_bb_alpha(
                bb1=bb1,
                bb2=bb2,
            ),
        ),
    )

complex_molecule = get_complex_molecule(
    bb1=stk.BuildingBlock('NCCN', [stk.PrimaryAminoFactory()]),
    bb2=stk.BuildingBlock('O=CC=O', [stk.AldehydeFactory()]),
)
other_complex_molecule = get_complex_molecule(
    bb1=stk.BuildingBlock('NCCCCCCCCN', [stk.PrimaryAminoFactory()]),
    bb2=stk.BuildingBlock('O=CCCCCCCCC=O', [stk.AldehydeFactory()]),
)

Anyway, as you say, one of the goals is extensibility so please do give feedback on your experience writing them, and feel free to ask questions.

[lukasturcani]

And if I wanted to keep info about the constituent building blocks, I would use composition rather than sublcassing here too:

class MyComplexMolecule:
    def __init__(self, bb1, bb2):
        self._stage1_constructed_molecule = stk.ConstructedMolecule(
            topology_graph=stk.SomeTopologyGraph(
                building_blocks=(bb1, bb2),
           ),
        )
        self._stage2_constructed_molecule = stk.ConstructedMolecule(
            topology_graph=stk.OtherTopologyGraph(
                bb1=stk.BuildingBlock.init_from_molecule(self._stage1_constructd_molecule),
            ),
        )

    def get_stage1_building_blocks():
        yield from self._stage1_constructed_molecule.get_building_blocks()

   def get_stage2_building_blocks():
       yield from self._stage2_constructed_molecule.get_building_blocks()

[joshkamm]

Thanks for the examples. That should be helpful for thinking about multiple layers of construction. I'm trying out a composition approach for the torsion wrapper API and it's been going reasonably well so far.

[joshkamm]

Thanks again for the help. I sent a PR over to stko: JelfsMaterialsGroup/stko#63

[andrewtarzia]

Hey @joshkamm,

This wrapper looks very interesting for analysis! We develop our analysis code into stko, as Lukas mentioned, but I also perform torsional analysis with my own written code and something like this could be super useful. In fact, we are interested in an extension of the stk molecule that is useful for optimisation processes, for which torsions are of interest!

Do you have an accessible repo for this code, otherwise, feel free to chuck a PR to the stko link Lukas sent and we can figure out how it fits into stko (at this stage stk would probably define a new stko.Molecule/BuildingBlock/ConstructedMolecule that inherits or contain the stk equivalent).

On this, I think the top two classes of your API would remain Molecule and ConstructedMolecule, but, unlike stk classes, these would now contain Atom/AtomInfo, Bond/BondInfo, Torsion/TorsionInfo (these can be calculated on the fly upon initialisation using rdkits path extraction code) and a position matrix?

[joshkamm]

In terms of use case, we are working on further developments related to this paper which our research group published last year applying machine learning to generate ensembles of low energy conformers:
https://papers.nips.cc/paper/2020/hash/e904831f48e729f9ad8355a894334700-Abstract.html
https://github.com/tarungog/torsionnet_paper_version

We hope to improve efficiency by decomposing macromolecules into subunits (e.g. oligomers of a polymer), training on the subunits (computationally cheap), and then transferring what was learned to reduce large molecule training (computationally expensive).

To do this our algorithmic learning agents would need access to an environment API that allows them to:

• Manipulate torsional angles of a macromolecule or any of its building blocks and receive feedback in the form of the energy of the corresponding locally optimized conformer.
• Map between a torsion in a macromolecule and the corresponding building block torsion

I've been using stk for molecular generation recently and Lukas added the underlying method I need for this current API in #267. In terms of stko specifically, we are currently using rdkit's MMFF optimizer for method development which already has a pretty simple interface. However, I expect in the future we will want something higher quality such as an XTB method to obtain chemically insightful results.

So are you saying I would implement MoleculeTorsioned and ConstructedMoleculeTorsioned as stk subclasses for now and then we could decide to merge the methods into stko specific molecule classes when I pull request?

[joshkamm]

Update: all the ways I found to convert a ConstructedMolecule generated through an stk construction process into a subclass seemed kind of ugly. This is one of the reasons I thought to ask if there was a better way to do the extension or if the functions would make sense in the stk library itself. It seems from Lukas's comments that he's expecting stk to be extended primarily by containment as well so I may play around with containment for this API to get some experience with what sorts of issues I run into.

[andrewtarzia]

Hey Josh, sorry for the delay!

That paper looks really interesting! I think if you implement that in stko, you have many energy methods available automatically, including xtb. So definitely, let us know how you get on. I am going to get started on an stko.Molecule class (or something like this) myself soon, which your class could merge into or contain. I would start with your own molecule that contains and manipulates an stk molecule (remembering to clone it and maintain immutability) in the ways you desire.

[andrewtarzia]

@lukasturcani and I were just chatting - and were wondering if the rdkit functions to set torsional angles were useful to your problem?

[joshkamm]

I forked the stko repository and have been committing progress on this torsion wrapping API to the fork. The files I've been working on are here: https://github.com/joshkamm/stko/tree/torsion-wrapper/stko/molecular/molecules and https://github.com/joshkamm/stko/tree/torsion-wrapper/stko/molecular/torsions .

I switched to cloning the stk molecule instead of assigning it in the wrapper constructor. I'm not aware if any of the methods I'm calling would modify the underlying stk molecule state. Do you mean that an stk molecule is supposed to be immutable or that a wrapper class that contains it should be immutable?

For your last comment, yes actually what we currently do is have our algorithmic agents generate an "action" which corresponds to a set of torsion angles. We call rdkit.Chem.rdMolTransforms.SetDihedralDeg() to set the torsions to the desired angles, perform a local optimization (currently rdkit MMFF), and give the agent feedback on this resulting new "state" of its action.

[andrewtarzia]

I would argue that this is answered by the new features in stko : JelfsMaterialsGroup/stko#96