Skip to content

Commit

Permalink
Merge pull request #35 from chrisjonesBSU/clean-up
Browse files Browse the repository at this point in the history 
Improve doc strings, fix remaining bugs left over after upgrading to hoomd4
  • Loading branch information
@chrisjonesBSU
chrisjonesBSU committed Feb 25, 2024
2 parents ecb2762 + cc50a4f commit 5c9d4c8
Show file tree
Hide file tree
Showing 3 changed files with 163 additions and 152 deletions.
156 changes: 99 additions & 57 deletions msibi/forces.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,10 @@
import warnings

from cmeutils.structure import (
angle_distribution, bond_distribution, dihedral_distribution, gsd_rdf
angle_distribution,
bond_distribution,
dihedral_distribution,
gsd_rdf
)
import matplotlib.pyplot as plt
import numpy as np
Expand All @@ -15,23 +18,44 @@


class Force(object):
"""Creates a potential, either to be held constant, or to be
optimized.
"""
Base class from which other forces inherit.
Don't call this class directly, instead use
msibi.forces.Bond, msibi.forces.Angle, msibi.forces.Pair,
and msibi.forces.Dihedral.
Forces in MSIBI can either be held constant (i.e. fixed) or
optimized (i.e. mutable). Only one type of of force
can be optimized at a time (i.e. angles, or pairs, etc..)
Parameters
----------
name : str, required
The name of the type in the bond.
Must match the names found in the State's .gsd trajectory file
The name of the type in the Force.
Must match the names found in the State's .gsd trajectory file.
optimize : bool, required
Set to True if this force is to be mutable and optimized.
Set to False if this force is to be held constant while
other forces are optimized.
nbins : int, optional
This must be a positive integer if this force is being optimized.
nbins is used to setting the potenials independent varible (x) range
and step size (dx).
"""

def __init__(
self,
name,
optimize=False,
optimize,
nbins=None,
head_correction_form="linear"
):
if optimize and nbins is None or nbins<=0:
raise ValueError(
"If a force is set to be optimized, nbins must be "
"a positive, non-zero integer."
)
self.name = name
self.optimize = optimize
self.head_correction_form = head_correction_form
Expand All @@ -42,29 +66,31 @@ def __init__(
self.x_range = None
self.potential_history = []
self._potential = None
self._potential_file = None
self._smoothing_window = 3
self._smoothing_order = 1
#TODO: set param for nbins?
self._nbins = nbins
self._force_type = None #TODO: Do we need this?
self._states = dict()
self._head_correction_history = []
self._tail_correction_history = []
self._learned_potential_history = []

if optimize and nbins is None:
raise ValueError(
"If this force is set to be optimized, the nbins "
"must be set as a non-zero value"
)

def __repr__(self):
return (
f"Type: {self.__class__}; "
+ f"Name: {self.name}; "
+ f"Optimize: {self.optimize}"
)

@property
def potential(self):
if self.format != "table":
#TODO: Set custom warning
warnings.warn(f"{self} is not using a table potential.")
return None
return self._potential

@potential.setter
Expand All @@ -82,7 +108,6 @@ def potential(self, array):
@property
def force(self):
if self.format != "table":
#TODO: Set custom warning
warnings.warn(f"{self} is not using a table potential.")
return None
return -1.0*np.gradient(self.potential, self.dx)
Expand Down Expand Up @@ -118,9 +143,30 @@ def nbins(self, value):
self._add_state(state)

def target_distribution(self, state):
"""The target structural distribution corresponding to this foce.
Parameters
----------
state : msibi.state.State, required
The state to use in finding the target distribution.
"""
return self._states[state]["target_distribution"]

def plot_target_distribution(self, state):
"""Quick plotting function that shows the target structural
distribution corresponding to this forces.
Parameters
----------
state : msibi.state.State, required
The state to use in finding the target distribution.
Notes
-----
Use this to see how the shape of the target distribution is
affected by your choices for nbins, smoothing window,
and smoothing order.
"""
#TODO: Make custom error
if not self.optimize:
raise RuntimeError(
Expand All @@ -142,26 +188,31 @@ def plot_target_distribution(self, state):
plt.plot(target[:,0], y_smoothed, label="Smoothed")
plt.legend()

def plot_potential_history(self):
if not self.optimize:
raise RuntimeError(
"This force object is not set to be optimized. "
)
fig = plt.figure()
for idx, i in enumerate(self.potential_history):
plt.plot(self.x_range, i, "o-", label=idx)
plt.legend(title="Iteration")
return fig

def plot_fit_scores(self, state):
if not self.optimize:
raise RuntimeError(
"This force object is not set to be optimized. "
)
fig = plt.figure()
plt.plot(self._states[state]["f_fit"], "o-")
plt.xlabel("Iteration")
plt.ylabel("Fit Score")
"""Returns a plot showing the evolution of the distribution
matching evolution.
Parameters
----------
state : msibi.state.State, required
The state to use in finding the target distribution.
"""
if not self.optimize:
raise RuntimeError("This force object is not set to be optimized.")
fig = plt.figure()
plt.plot(self._states[state]["f_fit"], "o-")
plt.xlabel("Iteration")
plt.ylabel("Fit Score")

def distribution_history(self, state):
"""Returns the complete query distribution history for a given state.
Parameters
----------
state : msibi.state.State, required
The state to use for calculating the distribution.
"""
return self._states[state]["distribution_history"]

def set_target_distribution(self, state, array):
""""""
Expand All @@ -178,7 +229,7 @@ def distribution_fit(self, state):
def set_quadratic(self, k4, k3, k2, x0, x_min, x_max):
"""Set a potential based on the following function:
V(x) = k4(l-x0)^4 + k3(l-x0)^3 + k2(l-x0)^2
V(x) = k4(x-x0)^4 + k3(x-x0)^3 + k2(x-x0)^2
Using this method will create a table potential V(x) over the range
x_min - x_max.
Expand All @@ -191,17 +242,15 @@ def set_quadratic(self, k4, k3, k2, x0, x_min, x_max):
x0, k4, k3, k2 : float, required
The paraters used in the V(x) function described above
x_min : float, required
The lower bound of the bond potential lengths
The lower bound of the potential range
x_max : float, required
The upper bound of the bond potential lengths
The upper bound of the potential range
"""
self.format = "table"
#TODO: Properties for these?
self.x_min = x_min
self.x_max = x_max
self.dx = x_max / self.nbins
self.x_range = np.arange(x_min, x_max+self.dx, self.dx)
self.x_range = np.arange(x_min, x_max + self.dx, self.dx)
self.potential = quadratic_spring(self.x_range, x0, k4, k3, k2)
self.force_init = "Table"
self.force_entry = self._table_entry()
Expand All @@ -213,25 +262,25 @@ def set_from_file(self, file_path):
is the potential enregy at r. The force will be calculated
from r and V using np.gradient().
Use this potential setter to set a potential from a previous MSIBI run.
For example, use the final potential files from a bond-optimization IBI
run to set a static coarse-grained bond potential while you perform
IBI runs on angle and pair potentials.
Parameters:
-----------
file_path : str, required
The full path to the table potential text file.
Notes
-----
Use this potential setter to set a potential from a previous MSIBI run.
For example, use the final potential files from a bond-optimization IBI
run to set a static coarse-grained bond potential while you perform
IBI runs on angle and/or pair potentials.
"""
self._potential_file = file_path
f = np.loadtxt(self._potential_file)
f = np.loadtxt(file_path)
self.x_range = f[:,0]
self.dx = np.round(self.x_range[1] - self.x_range[0], 3)
self.x_min = self.x_range[0]
self.x_max = self.x_range[-1] + self.dx
self._potential = f[:,1]
self.format = "table" #TODO: Still using format attribute?
self.format = "table"
self.force_init = "Table"
self.force_entry = self.table_entry()

Expand Down Expand Up @@ -263,11 +312,8 @@ def _add_state(self, state):
"target_distribution": target_distribution,
"current_distribution": None,
"alpha": state.alpha,
#TODO: Get rid of alpha form?
"alpha_form": "linear",
"f_fit": [],
"distribution_history": [],
#TODO: Is this used anywhere?
"path": state.dir
}

Expand Down Expand Up @@ -325,7 +371,6 @@ def _update_potential(self):
self.potential_history.append(np.copy(self.potential))
for state in self._states:
kT = state.kT
#TODO: Save distribution history?
current_dist = self._states[state]["current_distribution"]
target_dist = self._states[state]["target_distribution"]
self._states[state]["distribution_history"].append(current_dist)
Expand Down Expand Up @@ -356,7 +401,6 @@ def __init__(
self.type1, self.type2 = sorted(
[type1, type2], key=natural_sort
)
self._force_type = "bond"
self._correction_function = bond_correction
name = f"{self.type1}-{self.type2}"
super(Bond, self).__init__(
Expand Down Expand Up @@ -426,7 +470,7 @@ def __init__(
self.type2 = type2
self.type3 = type3
name = f"{self.type1}-{self.type2}-{self.type3}"
self._force_type = "angle"
self._correction_function = bond_correction
super(Angle, self).__init__(
name=name,
optimize=optimize,
Expand Down Expand Up @@ -461,17 +505,17 @@ def _table_entry(self):
table_entry = {"U": self.potential, "tau": self.force}
return table_entry

def _get_distribution(self, gsd_file):
def _get_distribution(self, state, gsd_file):
return angle_distribution(
gsd_file=gsd,
gsd_file=gsd_file,
A_name=self.type1,
B_name=self.type2,
C_name=self.type3,
start=-state.n_frames,
histogram=True,
normalize=True,
l_min=self.x_min,
l_max=self.x_max,
theta_min=self.x_min,
theta_max=self.x_max,
bins=self.nbins + 1
)

Expand All @@ -487,7 +531,6 @@ def __init__(
):
self.type1, self.type2 = sorted( [type1, type2], key=natural_sort)
name = f"{self.type1}-{self.type2}"
self._force_type = "pair"
self.r_cut = None
super(Pair, self).__init__(
name=name,
Expand Down Expand Up @@ -543,7 +586,6 @@ def __init__(
self.type3 = type3
self.type4 = type4
name = f"{self.type1}-{self.type2}-{self.type3}-{self.type4}"
self._force_type = "dihedral"
self.table_entry = dict(U=None, tau=None)
super(Dihedral, self).__init__(
name=name,
Expand Down
Loading

0 comments on commit 5c9d4c8

Please sign in to comment.