# AUTOGENERATED FILE! PLEASE DON'T EDIT
"""This module is for all things related to atoms, molecules and their simulations"""
import k1lib, torch, math, re
from functools import partial
from typing import Dict, List, Union
import matplotlib.pyplot as plt; import matplotlib as mpl
settings = {"overOctet": False}
__all__ = ["Atom", "System", "substances", "alkane", "endChain", "alcohol",
"parse", "sameEmpirical",
"settings", "NoFreeElectrons", "OctetFull"]
[docs]class NoFreeElectrons(RuntimeError): pass
[docs]class OctetFull(RuntimeError): pass
# if Atom's gDepth is smaller than this, then it means that it has not been visited
_depthAuto = k1lib.AutoIncrement()
_idxAuto = k1lib.AutoIncrement()
[docs]class Atom:
"""Just an atom really. Has properties, can bond to other atoms, and can
generate a :class:`System` for simulation."""
bonds: List["Atom"]
"""List of Atoms bonded to this Atom"""
[docs] def __init__(self, name:str, atomicN:int, massN:float, valenceE:int, radius:List[float]=[], octetE:int=8):
"""Creates a new atom. Not intended to be used by the end user. If you
wish to get a new atom, just do stuff like this::
c1 = mo.C
c2 = mo.C
c1 == c2 # returns False, demonstrating that these are different atoms
If you wish to register new substances with the module, you can do this::
genF = lambda: Atom(...)
mo.registerSubstance("elementName", genF)
mo.elementName # should executes `genF` and returns
:param name: element name (eg. "C")
:param atomicN: atomic number (eg. 6)
:param massN: atomic mass in g/mol (eg. 12)
:param valenceE: how many valence electrons initially?
:param radius: covalent radiuses (in pm) for single, double and triple bonds
:param octetE: how many electrons in a full octet? Default 8, but can be 2 for H and He"""
self.name = name; self.atomicN = atomicN; self.massN = massN
self.ogValenceE = valenceE # original
self.valenceE = valenceE; self.octetE = octetE; self.radius = radius
self.bonds = [] # list of Atoms this Atom is bonded to
self.gDepth = -1 # graph depth, for graph traversal stuff. Values will be updated from _depthAuto
self.idx = f"A{_idxAuto()}" # unique value for Atoms everywhere
# contracts:
# - valenceE = eClouds * 2 + freeE + len(bonds) * 2
# - valenceE <= octetE. "<" happens when octet not full
# can only form a new bond if freeE >= 1. Can dec eClouds to inc freeE
if name != "_e":
self.eClouds = []; self.freeE = valenceE % 2
for i in range(valenceE//2): self.eClouds.append(mo._e)
else: self.eClouds = []; self.freeE = 0
@property
def nonHBonds(self) -> List["Atom"]:
"""All bonds this atom has, minus the Hydrogens."""
return [a for a in self.bonds if a.name != "H"]
[docs] def nBonds(self, atom:"Atom"):
"""Get number of bonds between this and another atom."""
return len([bond for bond in self.bonds if bond == atom])
@property
def availableBonds(self) -> int:
"""Available bonds. This includes electron clouds, radical electrons, and
Hydrogen bonds."""
return len(self.eClouds) * 2 + self.freeE + len([a for a in self.bonds if a.name == "H"])
def __repr__(self):
return f"""<Atom {self.name} ({self.atomicN}), {len(self.bonds)} bonds, {self.valenceE}/{self.octetE} valence electrons, {len(self.eClouds)} electron clouds, {self.freeE} free (radical) electrons>"""
@k1lib.patch(Atom)
def _show(self, g=None, gDepth=-1, GVKwargs={}):
self.gDepth = gDepth; g.node(self.idx, self.name, **GVKwargs)
for atom in self.bonds:
if atom.gDepth >= gDepth: continue
d1 = (self.nonHBonds[0] == atom) if len(self.nonHBonds) > 0 else False
d2 = (atom.nonHBonds[0] == self) if len(atom.nonHBonds) > 0 else False
if d1 and d2: g(self.idx, atom.idx, dir="both")
elif d1: g(self.idx, atom.idx)
elif d2: g(atom.idx, self.idx)
else: g(self.idx, atom.idx, arrowhead="none")
[atom._show(g, gDepth) for atom in self.bonds if atom.gDepth < gDepth]
@k1lib.patch(Atom)
def show(self):
"""Show the molecule graph this atom is a part of. Meant for debugging
simple substances only, as graphs of big molecules look unwieldy. This also
highlights the current :class:`Atom`, and each bond is an arrow, indicating
where :meth:`next` will go next."""
g = k1lib.digraph(); self._show(g, _depthAuto(), {"style": "filled"}); return g
@k1lib.patch(Atom)
def _addFreeE(self, amt:int=1):
"""Adds free electron to atom."""
if amt > 1: [self._addFreeE() for i in range(amt)]
self.freeE += 1
if self.freeE >= 2: self.eClouds.append(mo._e); self.freeE -= 2
@k1lib.patch(Atom)
def _subFreeE(self, amt:int=1) -> bool:
"""Tries to use ``amt`` free electrons. Returns successful or not."""
if amt > 1: [self._subFreeE() for i in range(amt)]
elif self.freeE > 0: self.freeE -= 1
elif len(self.eClouds) > 0:
self.freeE += 1; self.eClouds.pop()
else: raise RuntimeError(f"Can't give away any more free electrons on atom {self.name}!")
@k1lib.patch(Atom)
def _makeRoom(self, nBonds:int):
"""Tries to remove bonds with Hydrogen to make room for ``nBonds`` more bonds."""
nBondsToRemove = self.valenceE + nBonds - self.octetE
if nBondsToRemove > 0:
Hs = [bond for bond in self.bonds if bond.name == "H"]
if len(Hs) >= nBondsToRemove:
for i in range(nBondsToRemove): self.removeBond(Hs[i])
elif not settings['overOctet']:
ans = input(f"Can't remove Hydrogen bonds to make room for new bond! Do you want to do anyway (y/n): ")
print("Btw, you can auto accept this by doing `mo.settings['overOctet'] = True`")
if ans.lower()[0] != "y": raise OctetFull("Stopping...")
availableE = len(self.eClouds) * 2 + self.freeE
if availableE < nBonds: raise NoFreeElectrons(f"Can't make room for {nBonds} new bonds on {self.name}. Only {availableE} electrons left for bonds!")
@k1lib.patch(Atom)
def __call__(self, atom:Atom, nBonds:int=1, main=False) -> Atom:
"""Forms a bond with another atom. If valence electrons are full, will
attempt to disconnect Hydrogens from self to make room.
:param bond: number of bonds. 2 for double, 3 for triple
:param main: whether to put this bond in front of existing bonds, to
signify the "main" chain, so that it works well with :meth:`next`
:return: self"""
self._makeRoom(nBonds); atom._makeRoom(nBonds)
if main: self.bonds = [atom] * nBonds + self.bonds
else: self.bonds += [atom] * nBonds
atom.bonds += [self] * nBonds
self.valenceE += nBonds; self._subFreeE(nBonds)
atom.valenceE += nBonds; atom._subFreeE(nBonds)
return self
@k1lib.patch(Atom)
def bond(self, atom:Atom, nBonds:int=1, main=False) -> Atom:
"""Like :meth:`__call__`, but returns the atom passed in instead, so you
can form the main loop quickly."""
self(atom, nBonds, main); return atom
@k1lib.patch(Atom)
def main(self, atom:Atom, nBonds:int=1) -> Atom:
"""Like :meth:`bond`, but with ``main`` param defaulted to True."""
return self.bond(atom, nBonds, True)
@k1lib.patch(Atom)
def removeBond(self, atom:"Atom"):
"""Removes all bonds between this and another atom"""
nBonds = self.nBonds(atom)
self.bonds = [bond for bond in self.bonds if bond != atom]
self.valenceE -= nBonds; self._addFreeE(nBonds)
atom.bonds = [bond for bond in atom.bonds if bond != self]
atom.valenceE -= nBonds; atom._addFreeE(nBonds)
@k1lib.patch(Atom, "next")
def _next(self, offset=0, times:int=1) -> "Atom":
"""Returns the next atom bonded to this. Tries to avoid going into Hydrogens.
This is the main way to navigate around the molecule.
You kinda have to make sure that your molecule's bonding order is appropriate by
choosing between :meth:`bond` and :meth:`main`. Check the bonding order with
:meth:`show`.
:param offset: if there are multiple non-Hydrogen atoms, which ones should I pick?
:param times: how many times do you want to chain ``.next()``?"""
if times < 0: raise RuntimeError("Can't do .next() with negative `times`")
if times == 0: return self
atoms = self.nonHBonds + [atom for atom in self.bonds if atom.name == "H"]
if len(atoms) == 0: return None
_next = atoms[offset]
if times == 1: return _next
else: return _next.next(offset, times-1)
@k1lib.patch(Atom)
def nexts(self, atoms:int=2) -> List[Atom]:
"""Kinda like :meth:`next`, but fetches multiple atoms on the backbone.
Example::
c1, c2 = mo.CH4(mo.CH4).nexts()"""
if atoms < 1: raise RuntimeError(f"Zero or negative ({atoms}) number of atoms does not make sense!")
if atoms == 1: return [self]
return [self, *(self.next().nexts(atoms-1))]
empiricalOrder = ["C", "H", "O", "N"]
def em1(e:str, n:int):
if n == 1: return e
else: return f"{e}{n}"
@k1lib.patch(Atom)
def _empirical(self, d:Dict[str, int], gDepth:int):
if self.gDepth >= gDepth: return
self.gDepth = gDepth; d[self.name] += 1
for atom in self.bonds: atom._empirical(d, gDepth)
@k1lib.patch(Atom)
def empirical(self) -> str:
"""Returns an empirical formula for the molecule this :class:`Atom` is attached to."""
d = k1lib.Object().withAutoDeclare(lambda: 0)
self._empirical(d, _depthAuto()); answer = ""
for e in empiricalOrder:
if e in d: answer += em1(e,d[e]); del d[e]
for e in d.state.keys(): answer += em1(e,d[e])
return answer
@k1lib.patch(Atom)
def _atoms(self, l, gDepth):
if self.gDepth >= gDepth: return
self.gDepth = gDepth; l.append(self)
for atom in self.bonds: atom._atoms(l, gDepth)
@k1lib.patch(Atom)
def atoms(self) -> List[Atom]:
"""Returns a list of Atoms in the molecule this specific Atom is attached to."""
l = []; self._atoms(l, _depthAuto()); return l
_a = {} # dict of atoms, which will be used to patch the entire module
class _Mo: pass
mo = _Mo() # convenience object so that I can use the same style as the module
def _atom(name, *args, **kwargs):
_a[name] = f = lambda: Atom(name, *args, **kwargs)
setattr(_Mo, name, property(partial((lambda self, f: f()), f=f)))
# covalent radius taken from (Pyykko & Atsumi) https://chem.libretexts.org/@api/deki/pages/2182/pdf/A3%253A%2bCovalent%2bRadii.pdf?stylesheet=default
_atom("_e", 0, 0.1, 0, [15]) # electron cloud, for internal use
_atom("H", 1, 1.008, 1, [32], octetE=2)
_atom("Li", 3, 6.94, 1, [133, 124])
_atom("Be", 4, 9.0122, 2, [102, 90, 85])
_atom("B", 5, 10.81, 3, [85, 78, 73])
_atom("C", 6, 12.011, 4, [75, 67, 60])
_atom("N", 7, 14.007, 5, [71, 60, 54])
_atom("O", 8, 15.999, 6, [63, 57, 53])
_atom("F", 9, 18.998, 7, [64, 59, 53])
_atom("Na", 11, 22.990, 1, [155, 160])
_atom("Mg", 12, 24.305, 2, [139, 132, 127])
_atom("Al", 13, 26.982, 3, [126, 113, 111])
_atom("Si", 14, 28.085, 4, [116, 107, 102])
_atom("P", 15, 30.974, 5, [111, 102, 94])
_atom("S", 16, 32.06, 6, [103, 94, 95])
_atom("Cl", 17, 35.45, 7, [99, 95, 93])
_atom("K", 19, 39.098, 1, [196, 193])
_atom("Ca", 20, 40.078, 2, [171, 147, 133])
_atom("Ga", 31, 69.723, 3, [124, 117, 121])
_atom("Ge", 32, 72.630, 4, [121, 111, 114])
_atom("As", 33, 74.922, 5, [121, 114, 106])
_atom("Se", 34, 78.971, 6, [116, 107, 107])
_atom("Br", 35, 79.904, 7, [114, 109, 110])
_atom("I", 53, 126.9, 7, [133, 129, 125])
def _mole(name, f):
_a[name] = f
setattr(_Mo, name, property(partial((lambda self, f: f()), f=f)))
_mole("H2O", lambda: mo.O(mo.H)(mo.H))
_mole("CH4", lambda: mo.C(mo.H)(mo.H)(mo.H)(mo.H))
_mole("COOH", lambda: mo.C(mo.O, 2)(mo.O(mo.H))(mo.H))
_mole("NH3", lambda: mo.N(mo.H)(mo.H)(mo.H))
_mole("CH3OH", lambda: mo.CH4(mo.H2O))
def glucose():
o = mo.O
o.bond(mo.CH4).bond(mo.CH3OH).bond(mo.CH3OH).bond(mo.CH3OH).bond(mo.CH3OH).bond(o)
o.next().bond(mo.CH3OH); return o
_mole("glucose", glucose)
def cyclohexane():
c = mo.CH4
return c.bond(mo.CH4).bond(mo.CH4).bond(mo.CH4).bond(mo.CH4).bond(mo.CH4).bond(c)
_mole("cyclohexane", cyclohexane)
def benzene():
c = mo.CH4
return c.main(mo.CH4, 2).main(mo.CH4).main(mo.CH4, 2).main(mo.CH4).main(mo.CH4, 2).main(c)
_mole("benzene", benzene)
def adenine():
n = mo.NH3
c1 = n.bond(mo.CH4).bond(mo.NH3, 2).bond(mo.CH4)
c2 = c1.bond(mo.CH4, 2); c2.bond(n)
c1.bond(mo.CH4)(mo.NH3).bond(mo.NH3, 2).bond(mo.CH4).bond(mo.NH3, 2).bond(c2);
return n
_mole("adenine", adenine)
def ribose():
o = mo.H2O
return o.bond(mo.CH4(mo.CH3OH)).main(mo.CH3OH).main(mo.CH3OH).main(mo.CH3OH).main(o)
_mole("ribose", ribose)
def adenosine(): ri = mo.ribose; ri.next(1)(mo.adenine); return ri
_mole("adenosine", adenosine)
[docs]def substances() -> List[str]:
"""Get a list of builtin substances. To register new substances, check over
:class:`Atom`."""
return [k for k in _a.keys() if not k.startswith("_")]
[docs]def alkane(n:int) -> Atom:
"""Creates an alkane with ``n`` carbons.
Example::
# returns "C3H8"
mo.alkane(3).empirical()"""
answer = mo.CH4
for i in range(n-1): answer = answer.bond(mo.CH4)
return answer
@k1lib.patch(Atom, "endChain")
@property
def endChain(a) -> Atom:
"""Do a bunch of .next() until reached the end of the carbon chain.
Example::
c1 = mo.alcohol(3, 1)
c3 = c1.endChain
c3(mo.NH3)
c1.show() # displays in cell"""
lastA = None
for i in range(200): # for loop to prevent infinite recursion
nextA = a.next()
if nextA == lastA: return a
lastA = a; a = nextA
[docs]def alcohol(n:int, loc:int=1) -> Atom:
"""Creates an alcohol with ``n`` carbons and an OH group at carbon ``loc``.
Example::
# returns "C3H8O"
mo.alcohol(3, 1).empirical()"""
a = alkane(n)
a.next(times=loc-1)(mo.H2O)
return a
@k1lib.patch(Atom)
def moveLastCTo2ndC(a:Atom) -> Atom:
"""Move last carbon to 2nd carbon. Useful in constructing iso- and tert-."""
end = endChain(a); nearEnd = end.next()
end.removeBond(nearEnd); nearEnd(mo.H); a.next()(mo.CH4); return a
@k1lib.patch(Atom)
def perfluoro_ize(mainA:Atom) -> Atom:
"""Replaces all C-H bonds with C-F bonds"""
for a in mainA.atoms():
b = a.next()
if a.name == "H" and b.name == "C":
a.removeBond(b); b(mo.F)
return mainA
def distV(x:torch.Tensor) -> torch.Tensor:
"""Distance vectors of points.
:param x: location Tensor of shape (n, 3)
:returns: vector Tensor of shape (n, n, 3)"""
n = x.shape[0]
return x.view(1, n, 3) - x.view(n, 1, 3)
@k1lib.patch(Atom)
def _s_bondLength(self, idx_i:Dict[str, int], bondLengths:torch.Tensor):
"""Calculates bond length for all bonds in this atom and stores in
``bondLengths``."""
bonds = list(set(self.bonds)) + self.eClouds
for atom in bonds:
nBonds = self.nBonds(atom)
bL = self.radius[nBonds-1] + atom.radius[nBonds-1]
bondLengths[idx_i[self.idx], idx_i[atom.idx]] = bL
bondLengths[idx_i[atom.idx], idx_i[self.idx]] = bL
[docs]class System:
[docs] def __init__(self, mapping:Dict[str, Atom]):
"""Creates a new system that contains a molecule so that it can be
simulated. Not intended to be used by the end user. Use :meth:`Atom.system`
instead.
:param mapping: maps from atom index to :class:`Atom` object"""
self.mapping = mapping; self.atoms = list(mapping.values())
self.i_idx = list(mapping.keys()); self.n = len(self.i_idx)
self.idx_i = {idx:i for i, idx in enumerate(self.i_idx)}
self.bondLengths = torch.zeros(self.n, self.n)
for atom in self.atoms: atom._s_bondLength(self.idx_i, self.bondLengths)
self.bondMask = (self.bondLengths > 0) + 0
self._x = 50*torch.randn(self.n, 3)
self._v = torch.zeros(self.n, 3)
self._a = torch.zeros(self.n, 3)
self.iMask = 1-torch.eye(self.n) # identity mask
def __getitem__(self, idx) -> Atom:
if isinstance(idx, str): return self.mapping[idx]
return self.mapping[self.i_idx[idx]]
[docs] def x(self, a:Atom) -> torch.Tensor:
"""Get current location vector of the specified :class:`Atom`."""
return self._x[self.idx_i[a.idx]]
[docs] def v(self, a:Atom):
"""Get current velocity vector of the specified :class:`Atom`."""
return self._v[self.idx_i[a.idx]]
@k1lib.patch(System)
def _calcForces(self, decay=2):
"""
:param decay: make this big to ignore distant atoms"""
n = self.n; dV = distV(self._x) # (n, n, 3)
l = torch.sqrt((dV**2).sum(dim=2)) + 0.001
dV = dV / l.view(n, n, 1) # normalized direction
aBond = 10 * dV * torch.tanh((l-self.bondLengths)*self.bondMask/100).view(n, n, 1)
vBond = aBond.sum(dim=1)
aCou = dV * ((100/l)**decay).view(n, n, 1)
vCou = .3 * aCou.clearNan().sum(dim=1)
#vCou *= .2/math.sqrt((vCou**2).sum(dim=1).max())
self._a = 10 * (vBond - vCou)
return aBond, aCou, l, dV, vBond, vCou
@k1lib.patch(System)
def _update(self, dt:float=1.0):
self._v += dt * self._a
self._v = torch.clamp(self._v * 0.8, -10, 10) # friction
self._x += dt * self._v
#self._x *= 0.99 # universe wants to quish things down
self._a.zero_()
@k1lib.patch(System)
def _changeDevice(self, device:str="cpu"):
self.bondLengths = self.bondLengths.to(device)
self.bondMask = self.bondMask.to(device)
self._x = self._x.to(device)
self._v = self._v.to(device)
self._a = self._a.to(device)
self.iMask = self.iMask.to(device)
@k1lib.patch(System)
def simulate(self, t:int=100, dt:float=1.0, recordXs:bool=True, cuda:bool=False) -> List[torch.Tensor]:
"""Simulate system for ``t`` steps.
:param t: simulation total steps
:param dt: simulation time between steps
:param recordXs: whether to record locations while the simulation happens. Put
False to save memory/performance.
:param cuda: if True, do the simulation on the graphics card
:return: if ``recordXs`` is True, returns max 100 location Tensors. The Tensors
will have shape of (n, 3), where n is the number of atoms and electron clouds
your molecule has."""
self._changeDevice("cuda" if cuda else "cpu")
every = (t // 100) or 1; xs = []
r1 = k1lib.Range(t); r2 = k1lib.Range(1, 2)
for i in range(t):
self._calcForces(decay=r1.toRange(r2, i)); self._update(dt)
if recordXs and i % every == 0:
xs.append(self._x.clone())
return xs
@k1lib.patch(Atom)
def _system(self, mapping:Dict[str, Atom], gDepth:int):
if self.gDepth >= gDepth: return
self.gDepth = gDepth
mapping[self.idx] = self
for atom in self.bonds: atom._system(mapping, gDepth)
for eCloud in self.eClouds: eCloud._system(mapping, gDepth)
@k1lib.patch(Atom)
def system(self) -> System:
"""Creates a :class:`System` of the molecule this :class:`Atom` is attached
to."""
mapping = dict(); self._system(mapping, _depthAuto()); return System(mapping)
@k1lib.patch(Atom)
def _plot(self, ax, x:torch.Tensor, idx_i:Dict[str, int]):
s = x[idx_i[self.idx]]
#bonds = set([atom for atom in self.bonds if atom.name != "_e"])
bonds = set([atom for atom in self.bonds])
for atom in bonds:
a = x[idx_i[atom.idx]]
ax.plot([s[0], a[0]], [s[1], a[1]], [s[2], a[2]])
@k1lib.patch(System)
def plot(self, x:torch.Tensor=None, ax:"matplotlib.axes.Axes"=None):
"""Plots the molecule.
Example::
%matplotlib widget
s = mo.CH4(mo.H2O).system()
s.simulate(); s.plot()
The first line is so that you can rotate the molecule around in 3d interactively.
The 3rd line includes a simulation step, to get the molecule into roughly the
right position.
:param x: location Tensor of shape (n, 3). If none provided, will use current locations
:param ax: Axes object, in case you want to plot this on an existing plot"""
if x is None: x = self._x
if ax is None: ax = plt.figure(dpi=150).add_subplot(projection="3d")
ax.clear(); com = x.mean(dim=0) # center of mass
box = math.sqrt(((x - com.view(1, -1))**2).sum(dim=1).max())
ax.set_xlim(com[0]-box, com[0]+box)
ax.set_ylim(com[1]-box, com[1]+box)
ax.set_zlim(com[2]-box, com[2]+box)
for atom in self.atoms:
if atom.name == "_e": continue
atom._plot(ax, x, self.idx_i)
s = x[self.idx_i[atom.idx]]
ax.text(s[0], s[1], s[2], f"{atom.name}", ha="center", va="center")
return ax
@k1lib.patch(System)
def animate(self, xs:List[torch.Tensor], rotateSpeed=0.5):
"""Animates the molecule. This requires location information from
:meth:`simulate`. Example::
s = mo.CH4(mo.H2O).system()
s.animate(s.simulate())"""
fig = plt.figure()
ax = fig.add_subplot(projection="3d")
plt.close(fig)
def update(frame):
self.plot(xs[frame], ax)
ax.view_init(azim=frame*rotateSpeed)
a = k1lib.viz.FAnim(fig, update, len(xs)); plt.close(fig); return a
empC = re.compile("[A-Z][a-z]*[0-9]*")
[docs]def sameEmpirical(a:str, b:str) -> bool:
"""Checks whether 2 empirical formula are the same or not.
Example::
moparse.sameEmpirical("C2H4", "H4C2") # returns True"""
return sorted(re.findall(empC, a)) == sorted(re.findall(empC, b))
[docs]def parse(s:str, quiet:bool=False) -> Union[Atom, str]:
"""Tries to recognize molecule. Returns molecule :class:`~k1lib.mo.Atom`
if recognized, else the molecule's string."""
return k1lib._moparse.recognize(s, quiet)