diff --git a/examples/classical/gbForce/10p.pdb b/examples/classical/gbForce/10p.pdb
new file mode 100644
index 000000000..05b18aa71
--- /dev/null
+++ b/examples/classical/gbForce/10p.pdb
@@ -0,0 +1,25 @@
+TITLE GRoups of Organic Molecules in ACtion for Science
+REMARK THIS IS A SIMULATION BOX
+CRYST1 30.000 30.000 30.000 90.00 90.00 90.00 P 1 1
+MODEL 1
+ATOM 1 TP1 TE1 A 1 28.725 22.974 33.406 1.00 0.00 P
+ATOM 2 P00 IN1 A 2 26.892 23.631 31.811 1.00 0.00 C
+ATOM 3 P00 IN1 A 3 27.633 24.222 29.541 1.00 0.00 C
+ATOM 4 P00 IN1 A 4 25.695 23.995 27.884 1.00 0.00 C
+ATOM 5 P00 IN1 A 5 26.290 25.095 25.743 1.00 0.00 C
+ATOM 6 P00 IN1 A 6 24.826 24.484 23.700 1.00 0.00 C
+ATOM 7 P00 IN1 A 7 24.944 26.058 21.774 1.00 0.00 C
+ATOM 8 P00 IN1 A 8 23.018 26.878 20.449 1.00 0.00 C
+ATOM 9 P00 IN1 A 9 22.109 25.856 18.341 1.00 0.00 C
+ATOM 10 TP1 TE1 A 10 19.868 26.802 17.355 1.00 0.00 P
+TER
+CONECT 1 2
+CONECT 2 3
+CONECT 3 4
+CONECT 4 5
+CONECT 5 6
+CONECT 6 7
+CONECT 7 8
+CONECT 8 9
+CONECT 9 10
+ENDMDL
diff --git a/examples/classical/gbForce/1_5corrV2.xml b/examples/classical/gbForce/1_5corrV2.xml
new file mode 100644
index 000000000..633b4326b
--- /dev/null
+++ b/examples/classical/gbForce/1_5corrV2.xml
@@ -0,0 +1,98 @@
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+ step(r+sr2-or1)*0.5*(1/L-1/U+0.25*(1/U^2-1/L^2)*(r-sr2*sr2/r)+0.5*log(L/U)/r+C);
+ U=r+sr2; C=2*(1/or1-1/L)*step(sr2-r-or1); L=max(or1, D); D=abs(r-sr2); sr2 =
+ scale2*or2; or1 = radius1-0.009; or2 = radius2-0.009
+
+
+ 1/(1/or-tanh(1*psi-0.8*psi^2+4.85*psi^3)/radius); psi=I*or; or=radius-0.009
+
+
+ 28.3919551*(radius+0.14)^2*(radius/B)^6-0.5*138.935456*(1/soluteDielectric-1/solventDielectric)*charge^2/B
+
+
+ -138.935456*(1/soluteDielectric-1/solventDielectric)*charge1*charge2/f;
+ f=sqrt(r^2+B1*B2*exp(-r^2/(4*B1*B2)))
+
+
+
+
+
+
+
+
+
+
+
+
+
+
\ No newline at end of file
diff --git a/examples/classical/gbForce/gbforce.py b/examples/classical/gbForce/gbforce.py
new file mode 100644
index 000000000..55a42756f
--- /dev/null
+++ b/examples/classical/gbForce/gbforce.py
@@ -0,0 +1,99 @@
+#!/usr/bin/env python
+from openmm import *
+from openmm.app import *
+from openmm.unit import *
+import numpy as np
+import sys
+sys.path.append("..")
+sys.path.append(".")
+sys.path.append("...")
+from dmff.api.hamiltonian import Hamiltonian
+from dmff.common import nblist
+from jax import jit
+import jax.numpy as jnp
+import mdtraj as md
+
+def forcegroupify(system):
+ forcegroups = {}
+ for i in range(system.getNumForces()):
+ force = system.getForce(i)
+ force.setForceGroup(i)
+ forcegroups[force] = i
+ return forcegroups
+
+def getEnergyDecomposition(context, forcegroups):
+ energies = {}
+ for f, i in forcegroups.items():
+ energies[f] = context.getState(getEnergy=True, groups=2 ** i).getPotentialEnergy()
+ return energies
+
+print("MM Reference Energy:")
+pdb = PDBFile("10p.pdb")
+ff = ForceField("1_5corrV2.xml")
+system = ff.createSystem(pdb.topology, nonbondedMethod=NoCutoff, constraints=None, removeCMMotion=False)
+h = Hamiltonian("1_5corrV2.xml")
+params = h.getParameters()
+compoundBondForceParam = params["Custom1_5BondForce"]
+length = compoundBondForceParam["length"]
+k = compoundBondForceParam["k"]
+customCompoundForce = openmm.CustomCompoundBondForce(2, "0.5*k*(distance(p1,p2)-length)^2")
+customCompoundForce.addPerBondParameter("length")
+customCompoundForce.addPerBondParameter("k")
+for i, leng in enumerate(length):
+ customCompoundForce.addBond([i, i+4], [leng, k[i]])
+system.addForce(customCompoundForce)
+print("Dih info:")
+for force in system.getForces():
+ if isinstance(force, PeriodicTorsionForce):
+ print("No. of dihs:", force.getNumTorsions())
+
+forcegroups = forcegroupify(system)
+integrator = VerletIntegrator(0.1)
+context = Context(system, integrator, Platform.getPlatformByName("Reference"))
+context.setPositions(pdb.positions)
+state = context.getState(getEnergy=True)
+energy = state.getPotentialEnergy()
+energies = getEnergyDecomposition(context, forcegroups)
+print("Total energy:", energy)
+for key in energies.keys():
+ print(key.getName(), energies[key])
+
+print("Jax Energy")
+h = Hamiltonian("1_5corrV2.xml")
+pot = h.createPotential(pdb.topology, nonbondedMethod=NoCutoff)
+params = h.getParameters()
+positions = pdb.getPositions(asNumpy=True).value_in_unit(nanometer)
+positions = jnp.array(positions)
+box = np.array([
+ [30.0, 0.0, 0.0],
+ [0.0, 30.0, 0.0],
+ [0.0, 0.0, 30.0]
+])
+
+# neighbor list
+rc = 6.0
+nbl = nblist.NeighborList(box, rc, pot.meta['cov_map'])
+nbl.allocate(positions)
+pairs = nbl.pairs
+
+bondE = pot.dmff_potentials['HarmonicBondForce']
+print("Bond:", bondE(positions, box, pairs, params))
+
+angleE = pot.dmff_potentials['HarmonicAngleForce']
+print("Angle:", angleE(positions, box, pairs, params))
+
+gbE = pot.dmff_potentials['CustomGBForce']
+print("CustomGBForce:", gbE(positions, box, pairs, params))
+
+E1_5 = pot.dmff_potentials['Custom1_5BondForce']
+print("Custom1_5BondForce:", E1_5(positions, box, pairs, params))
+
+dihE = pot.dmff_potentials['CustomTorsionForce']
+print("Torsion:", dihE(positions, box, pairs, params))
+
+nbE = pot.dmff_potentials['NonbondedForce']
+print("Nonbonded:", nbE(positions, box, pairs, params))
+
+etotal = pot.getPotentialFunc()
+print("Total:", etotal(positions, box, pairs, params))
+
diff --git a/tests/data/10p.pdb b/tests/data/10p.pdb
new file mode 100644
index 000000000..05b18aa71
--- /dev/null
+++ b/tests/data/10p.pdb
@@ -0,0 +1,25 @@
+TITLE GRoups of Organic Molecules in ACtion for Science
+REMARK THIS IS A SIMULATION BOX
+CRYST1 30.000 30.000 30.000 90.00 90.00 90.00 P 1 1
+MODEL 1
+ATOM 1 TP1 TE1 A 1 28.725 22.974 33.406 1.00 0.00 P
+ATOM 2 P00 IN1 A 2 26.892 23.631 31.811 1.00 0.00 C
+ATOM 3 P00 IN1 A 3 27.633 24.222 29.541 1.00 0.00 C
+ATOM 4 P00 IN1 A 4 25.695 23.995 27.884 1.00 0.00 C
+ATOM 5 P00 IN1 A 5 26.290 25.095 25.743 1.00 0.00 C
+ATOM 6 P00 IN1 A 6 24.826 24.484 23.700 1.00 0.00 C
+ATOM 7 P00 IN1 A 7 24.944 26.058 21.774 1.00 0.00 C
+ATOM 8 P00 IN1 A 8 23.018 26.878 20.449 1.00 0.00 C
+ATOM 9 P00 IN1 A 9 22.109 25.856 18.341 1.00 0.00 C
+ATOM 10 TP1 TE1 A 10 19.868 26.802 17.355 1.00 0.00 P
+TER
+CONECT 1 2
+CONECT 2 3
+CONECT 3 4
+CONECT 4 5
+CONECT 5 6
+CONECT 6 7
+CONECT 7 8
+CONECT 8 9
+CONECT 9 10
+ENDMDL
diff --git a/tests/data/1_5corrV2.xml b/tests/data/1_5corrV2.xml
new file mode 100644
index 000000000..633b4326b
--- /dev/null
+++ b/tests/data/1_5corrV2.xml
@@ -0,0 +1,98 @@
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+ step(r+sr2-or1)*0.5*(1/L-1/U+0.25*(1/U^2-1/L^2)*(r-sr2*sr2/r)+0.5*log(L/U)/r+C);
+ U=r+sr2; C=2*(1/or1-1/L)*step(sr2-r-or1); L=max(or1, D); D=abs(r-sr2); sr2 =
+ scale2*or2; or1 = radius1-0.009; or2 = radius2-0.009
+
+
+ 1/(1/or-tanh(1*psi-0.8*psi^2+4.85*psi^3)/radius); psi=I*or; or=radius-0.009
+
+
+ 28.3919551*(radius+0.14)^2*(radius/B)^6-0.5*138.935456*(1/soluteDielectric-1/solventDielectric)*charge^2/B
+
+
+ -138.935456*(1/soluteDielectric-1/solventDielectric)*charge1*charge2/f;
+ f=sqrt(r^2+B1*B2*exp(-r^2/(4*B1*B2)))
+
+
+
+
+
+
+
+
+
+
+
+
+
+
\ No newline at end of file
diff --git a/tests/data/pBox.pdb b/tests/data/pBox.pdb
new file mode 100644
index 000000000..b1a0118e1
--- /dev/null
+++ b/tests/data/pBox.pdb
@@ -0,0 +1,87 @@
+TITLE polyten_GMX.gro created by acpype (v: 2022.6.6) on Wed Oct 5 08:03:17 2022
+REMARK THIS IS A SIMULATION BOX
+CRYST1 30.000 30.000 30.000 90.00 90.00 90.00 P 1 1
+MODEL 1
+HETATM 1 P00 TE1 1 28.652 12.361 14.349 1.00 0.00 P
+HETATM 2 O01 TE1 1 28.511 11.590 12.887 1.00 0.00 O
+HETATM 3 O02 TE1 1 28.857 11.267 15.580 1.00 0.00 O
+HETATM 4 O03 TE1 1 29.865 13.487 14.329 1.00 0.00 O
+HETATM 5 O04 TE1 1 26.965 13.160 14.653 1.00 0.00 O
+HETATM 6 P00 IN1 2 26.260 14.805 14.965 1.00 0.00 P
+HETATM 7 O01 IN1 2 26.855 15.339 16.380 1.00 0.00 O
+HETATM 8 O02 IN1 2 26.508 15.706 13.631 1.00 0.00 O
+HETATM 9 O00 IN2 3 24.431 14.542 15.100 1.00 0.00 O
+HETATM 10 P00 IN1 4 23.086 15.420 16.142 1.00 0.00 P
+HETATM 11 O01 IN1 4 22.975 14.492 17.468 1.00 0.00 O
+HETATM 12 O02 IN1 4 23.530 16.970 16.231 1.00 0.00 O
+HETATM 13 O00 IN2 5 21.421 15.307 15.326 1.00 0.00 O
+HETATM 14 P00 IN1 6 19.909 16.539 15.340 1.00 0.00 P
+HETATM 15 O01 IN1 6 20.000 17.324 16.743 1.00 0.00 O
+HETATM 16 O02 IN1 6 20.063 17.265 13.902 1.00 0.00 O
+HETATM 17 O00 IN2 7 18.227 15.697 15.269 1.00 0.00 O
+HETATM 18 P00 IN1 8 16.526 16.321 15.943 1.00 0.00 P
+HETATM 19 O01 IN1 8 16.373 15.511 17.337 1.00 0.00 O
+HETATM 20 O02 IN1 8 16.580 17.927 15.852 1.00 0.00 O
+HETATM 21 O00 IN2 9 15.002 15.739 14.963 1.00 0.00 O
+HETATM 22 P00 IN1 10 13.371 16.585 14.465 1.00 0.00 P
+HETATM 23 O01 IN1 10 13.092 17.763 15.532 1.00 0.00 O
+HETATM 24 O02 IN1 10 13.581 16.841 12.885 1.00 0.00 O
+HETATM 25 O00 IN2 11 11.772 15.515 14.643 1.00 0.00 O
+HETATM 26 P00 IN1 12 10.263 15.275 13.556 1.00 0.00 P
+HETATM 27 O01 IN1 12 10.068 16.593 12.650 1.00 0.00 O
+HETATM 28 O02 IN1 12 10.482 13.800 12.922 1.00 0.00 O
+HETATM 29 O00 IN2 13 8.598 15.069 14.540 1.00 0.00 O
+HETATM 30 P00 IN1 14 6.866 15.556 14.061 1.00 0.00 P
+HETATM 31 O01 IN1 14 6.615 17.010 14.733 1.00 0.00 O
+HETATM 32 O02 IN1 14 6.677 15.310 12.476 1.00 0.00 O
+HETATM 33 O00 IN2 15 5.578 14.441 14.933 1.00 0.00 O
+HETATM 34 P00 IN1 16 3.855 13.938 14.474 1.00 0.00 P
+HETATM 35 O01 IN1 16 3.150 15.098 13.579 1.00 0.00 O
+HETATM 36 O02 IN1 16 3.961 12.437 13.850 1.00 0.00 O
+HETATM 37 O04 TE1 17 2.946 13.817 16.041 1.00 0.00 O
+HETATM 38 P00 TE1 17 1.214 13.377 16.663 1.00 0.00 P
+HETATM 39 O01 TE1 17 0.217 12.929 15.419 1.00 0.00 O
+HETATM 40 O02 TE1 17 0.655 14.736 17.433 1.00 0.00 O
+HETATM 41 O03 TE1 17 1.442 12.135 17.739 1.00 0.00 O
+TER
+CONECT 1 2
+CONECT 1 3
+CONECT 1 4
+CONECT 1 5
+CONECT 5 6
+CONECT 6 7
+CONECT 6 8
+CONECT 6 9
+CONECT 9 10
+CONECT 10 11
+CONECT 10 12
+CONECT 10 13
+CONECT 13 14
+CONECT 14 15
+CONECT 14 16
+CONECT 14 17
+CONECT 17 18
+CONECT 18 19
+CONECT 18 20
+CONECT 18 21
+CONECT 21 22
+CONECT 22 23
+CONECT 22 24
+CONECT 22 25
+CONECT 25 26
+CONECT 26 27
+CONECT 26 28
+CONECT 26 29
+CONECT 29 30
+CONECT 30 31
+CONECT 30 32
+CONECT 30 33
+CONECT 33 34
+CONECT 34 35
+CONECT 34 36
+CONECT 34 37
+CONECT 37 38
+CONECT 38 39
+CONECT 38 40
+CONECT 38 41
+ENDMDL
diff --git a/tests/data/polyp_amberImp.xml b/tests/data/polyp_amberImp.xml
new file mode 100644
index 000000000..15f6c9721
--- /dev/null
+++ b/tests/data/polyp_amberImp.xml
@@ -0,0 +1,111 @@
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+ step(r+sr2-or1)*0.5*(1/L-1/U+0.25*(1/U^2-1/L^2)*(r-sr2*sr2/r)+0.5*log(L/U)/r+C);
+ U=r+sr2; C=2*(1/or1-1/L)*step(sr2-r-or1); L=max(or1, D); D=abs(r-sr2); sr2 =
+ scale2*or2; or1 = radius1-0.009; or2 = radius2-0.009
+
+
+ 1/(1/or-tanh(1*psi-0.8*psi^2+4.85*psi^3)/radius); psi=I*or; or=radius-0.009
+
+
+ 28.3919551*(radius+0.14)^2*(radius/B)^6-0.5*138.935456*
+ (1/soluteDielectric-1/solventDielectric)*charge^2/B
+
+
+ -138.935456*(1/soluteDielectric-1/solventDielectric)*charge1*charge2/f;
+ f=sqrt(r^2+B1*B2*exp(-r^2/(4*B1*B2)))
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
\ No newline at end of file
diff --git a/tests/test_classical/test_gbforce.py b/tests/test_classical/test_gbforce.py
new file mode 100644
index 000000000..ffa29c5a7
--- /dev/null
+++ b/tests/test_classical/test_gbforce.py
@@ -0,0 +1,80 @@
+import pytest
+import jax
+import jax.numpy as jnp
+import openmm.app as app
+import openmm.unit as unit
+import numpy as np
+import numpy.testing as npt
+from dmff.api import Hamiltonian
+from dmff.common import nblist
+
+
+@pytest.mark.parametrize(
+ "pdb, prm, value",
+ [
+ ("./tests/data/10p.pdb", "./tests/data/1_5corrV2.xml", -11184.921239189738),
+ ("./tests/data/pBox.pdb", "./tests/data/polyp_amberImp.xml", -13914.34177591779),
+ ])
+def test_custom_gb_force(pdb, prm, value):
+ pdb = app.PDBFile(pdb)
+ h = Hamiltonian(prm)
+ potential = h.createPotential(
+ pdb.topology,
+ nonbondedMethod=app.NoCutoff
+ )
+ pos = jnp.asarray(pdb.getPositions(asNumpy=True).value_in_unit(unit.nanometer))
+ box = np.array([[20.0, 0.0, 0.0], [0.0, 20.0, 0.0], [0.0, 0.0, 20.0]])
+ rc = 6.0
+ nbl = nblist.NeighborList(box, rc, potential.meta['cov_map'])
+ nbl.allocate(pos)
+ pairs = nbl.pairs
+ gbE = potential.getPotentialFunc(names=["CustomGBForce"])
+ energy = gbE(pos, box, pairs, h.paramset)
+ print(energy)
+ npt.assert_almost_equal(energy, value, decimal=3)
+
+
+@pytest.mark.parametrize(
+ "pdb, prm, value",
+ [
+ ("./tests/data/10p.pdb", "./tests/data/1_5corrV2.xml", 59.53033875302844),
+ ])
+def test_custom_torsion_force(pdb, prm, value):
+ pdb = app.PDBFile(pdb)
+ h = Hamiltonian(prm)
+ potential = h.createPotential(
+ pdb.topology,
+ nonbondedMethod=app.NoCutoff
+ )
+ pos = jnp.asarray(pdb.getPositions(asNumpy=True).value_in_unit(unit.nanometer))
+ box = np.array([[20.0, 0.0, 0.0], [0.0, 20.0, 0.0], [0.0, 0.0, 20.0]])
+ rc = 6.0
+ nbl = nblist.NeighborList(box, rc, potential.meta['cov_map'])
+ nbl.allocate(pos)
+ pairs = nbl.pairs
+ gbE = potential.getPotentialFunc(names=["CustomTorsionForce"])
+ energy = gbE(pos, box, pairs, h.paramset)
+ npt.assert_almost_equal(energy, value, decimal=3)
+
+
+@pytest.mark.parametrize(
+ "pdb, prm, value",
+ [
+ ("./tests/data/10p.pdb", "./tests/data/1_5corrV2.xml", 117.95416362791674),
+ ])
+def test_custom_1_5bond_force(pdb, prm, value):
+ pdb = app.PDBFile(pdb)
+ h = Hamiltonian(prm)
+ potential = h.createPotential(
+ pdb.topology,
+ nonbondedMethod=app.NoCutoff
+ )
+ pos = jnp.asarray(pdb.getPositions(asNumpy=True).value_in_unit(unit.nanometer))
+ box = np.array([[20.0, 0.0, 0.0], [0.0, 20.0, 0.0], [0.0, 0.0, 20.0]])
+ rc = 6.0
+ nbl = nblist.NeighborList(box, rc, potential.meta['cov_map'])
+ nbl.allocate(pos)
+ pairs = nbl.pairs
+ gbE = potential.getPotentialFunc(names=["Custom1_5BondForce"])
+ energy = gbE(pos, box, pairs, h.paramset)
+ npt.assert_almost_equal(energy, value, decimal=3)
\ No newline at end of file