diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 7b1e4f9..3fc41fd 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -9,7 +9,7 @@ ci: submodules: false repos: - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.15.22 # Ruff version + rev: v0.16.6 # Ruff version hooks: - id: ruff args: [--fix, --extend-ignore=E203] diff --git a/examples/grits.ipynb b/examples/grits.ipynb index 8b3696c..c8db65a 100644 --- a/examples/grits.ipynb +++ b/examples/grits.ipynb @@ -21,8 +21,7 @@ "source": [ "import mbuild as mb\n", "\n", - "from grits import backmap\n", - "from grits import CG_Compound" + "from grits import CG_Compound, backmap" ] }, { diff --git a/grits/__init__.py b/grits/__init__.py index bdcac03..ec7e0b9 100644 --- a/grits/__init__.py +++ b/grits/__init__.py @@ -13,10 +13,10 @@ pass __all__ = [ - "__version__", + "Bead", "CG_Compound", "CG_System", - "Bead", + "__version__", "backmap", "utils", ] diff --git a/grits/coarsegrain.py b/grits/coarsegrain.py index efb6fdf..3561315 100644 --- a/grits/coarsegrain.py +++ b/grits/coarsegrain.py @@ -29,7 +29,7 @@ has_number, ) -__all__ = ["CG_Compound", "CG_System", "Bead"] +__all__ = ["Bead", "CG_Compound", "CG_System"] class CG_Compound(Compound): @@ -101,7 +101,7 @@ def __init__( aniso_beads=False, **kwargs, ): - super(CG_Compound, self).__init__(**kwargs) + super().__init__(**kwargs) if (beads is None) == (mapping is None): raise ValueError( "Please provide only one of either beads or mapping." @@ -120,7 +120,7 @@ def __init__( # to be set correctly. with tempfile.NamedTemporaryFile() as f: mol.write(format="mol2", filename=f.name, overwrite=True) - mol = list(pybel.readfile("mol2", f.name))[0] + mol = next(iter(pybel.readfile("mol2", f.name))) mol.OBMol.AddHydrogens() # mol.addh() n_atoms2 = mol.OBMol.NumAtoms() @@ -294,7 +294,11 @@ def save_mapping(self, filename=None): return filename def visualize( - self, show_ports=False, color_scheme={}, show_atomistic=False, scale=1.0 + self, + show_ports=False, + color_scheme=None, + show_atomistic=False, + scale=1.0, ): # pragma: no cover """Visualize the Compound using py3dmol. @@ -322,6 +326,8 @@ def visualize( ------- view : py3Dmol.view """ + if color_scheme is None: + color_scheme = {} if not run_from_ipython(): raise RuntimeError( "Visualization is only supported in Jupyter Notebooks." @@ -455,7 +461,7 @@ class Bead(Compound): def __init__(self, smarts=None, orientation=None, **kwargs): self.smarts = smarts self.orientation = orientation - super(Bead, self).__init__(element=None, **kwargs) + super().__init__(element=None, **kwargs) class CG_System: @@ -752,7 +758,7 @@ def save(self, cg_gsdfile, start=0, stop=None, stride=1): heavy_positions = positions[ np.where(masses > hmass) ] - major_axis, ab_idxs = get_major_axis( + major_axis, _ab_idxs = get_major_axis( heavy_positions ) orientation.append(get_quaternion(major_axis)) diff --git a/grits/finegrain.py b/grits/finegrain.py index fc0f0e7..52b7799 100644 --- a/grits/finegrain.py +++ b/grits/finegrain.py @@ -29,12 +29,12 @@ def fg_particles(): """Set the particles of the fine-grained structure.""" fine_grained = Compound() - anchors = dict() + anchors = {} for i, bead in enumerate(cg_compound): smiles = bead.smarts b = load(smiles, smiles=True) b.translate_to(bead.pos) - anchors[i] = dict() + anchors[i] = {} if cg_compound.anchors is not None: for index in cg_compound.anchors[bead.name]: anchors[i][index] = b[index] @@ -45,7 +45,7 @@ def fg_bonds(): """Set the bonds for the fine-grained structure.""" bonded_atoms = [] remove_hs = [] - rotated = {k: False for k in anchors.keys()} + rotated = {k: False for k in anchors} for name, inds in cg_compound.bond_map: for ibead, jbead in cg_compound.bonds(): names = [ibead.name, jbead.name] @@ -61,12 +61,12 @@ def fg_bonds(): try: iatom = anchors[i].pop(fi) except KeyError: - fi = [x for x in inds if x in anchors[i]][0] + fi = next(x for x in inds if x in anchors[i]) iatom = anchors[i].pop(fi) try: jatom = anchors[j].pop(fj) except KeyError: - fj = [x for x in inds if x in anchors[j]][0] + fj = next(x for x in inds if x in anchors[j]) jatom = anchors[j].pop(fj) hi = get_hydrogen(fine_grained, iatom) diff --git a/grits/tests/test_coarsegrain.py b/grits/tests/test_coarsegrain.py index 424cfe6..bec57f4 100644 --- a/grits/tests/test_coarsegrain.py +++ b/grits/tests/test_coarsegrain.py @@ -25,7 +25,7 @@ def test_initmethane(self, methane): assert cg_methane.n_particles == 1 assert isinstance(cg_methane, CG_Compound) - types = set([i.name for i in cg_methane.particles()]) + types = {i.name for i in cg_methane.particles()} assert "_A" in types assert len(types) == 1 assert np.isclose(cg_methane.mass, 16.043, atol=1e-5) @@ -38,7 +38,7 @@ def test_initanisomethane(self, methane): assert cg_methane.n_particles == 1 assert isinstance(cg_methane, CG_Compound) - types = set([i.name for i in cg_methane.particles()]) + types = {i.name for i in cg_methane.particles()} assert "_A" in types assert len(types) == 1 assert np.isclose(cg_methane.mass, 16.043, atol=1e-5) @@ -54,7 +54,7 @@ def test_initp3ht(self, p3ht): assert cg_p3ht.n_particles == 48 assert isinstance(cg_p3ht, CG_Compound) - types = set([i.name for i in cg_p3ht.particles()]) + types = {i.name for i in cg_p3ht.particles()} assert "_B" in types assert "_S" in types assert len(types) == 2 @@ -69,7 +69,7 @@ def test_initanisop3ht(self, p3ht): assert cg_p3ht.n_particles == 48 assert isinstance(cg_p3ht, CG_Compound) - types = set([i.name for i in cg_p3ht.particles()]) + types = {i.name for i in cg_p3ht.particles()} assert "_B" in types assert "_S" in types assert len(types) == 2 @@ -87,7 +87,7 @@ def test_initp3htoverlap(self, p3ht): assert cg_p3ht.n_particles == 48 assert isinstance(cg_p3ht, CG_Compound) - types = set([i.name for i in cg_p3ht.particles()]) + types = {i.name for i in cg_p3ht.particles()} assert "_B" in types assert "_S" in types assert len(types) == 2 @@ -102,7 +102,7 @@ def test_initanisop3htoverlap(self, p3ht): assert cg_p3ht.n_particles == 48 assert isinstance(cg_p3ht, CG_Compound) - types = set([i.name for i in cg_p3ht.particles()]) + types = {i.name for i in cg_p3ht.particles()} assert "_B" in types assert "_S" in types assert len(types) == 2 @@ -115,7 +115,7 @@ def test_initmapp3ht(self, p3ht, p3ht_mapping): assert cg_p3ht.n_particles == 48 assert isinstance(cg_p3ht, CG_Compound) - types = set([i.name for i in cg_p3ht.particles()]) + types = {i.name for i in cg_p3ht.particles()} assert "_B" in types assert "_S" in types assert len(types) == 2 @@ -126,7 +126,7 @@ def test_initmapanisop3ht(self, p3ht, p3ht_mapping): assert cg_p3ht.n_particles == 48 assert isinstance(cg_p3ht, CG_Compound) - types = set([i.name for i in cg_p3ht.particles()]) + types = {i.name for i in cg_p3ht.particles()} assert "_B" in types assert "_S" in types assert len(types) == 2 @@ -139,7 +139,7 @@ def test_initmapmethane(self, methane, methane_mapping): assert cg_methane.n_particles == 1 assert isinstance(cg_methane, CG_Compound) - types = set([i.name for i in cg_methane.particles()]) + types = {i.name for i in cg_methane.particles()} assert "_A" in types assert len(types) == 1 @@ -151,7 +151,7 @@ def test_initmapanisomethane(self, methane, methane_mapping): assert cg_methane.n_particles == 1 assert isinstance(cg_methane, CG_Compound) - types = set([i.name for i in cg_methane.particles()]) + types = {i.name for i in cg_methane.particles()} assert "_A" in types assert len(types) == 1 for particle in cg_methane.particles(): diff --git a/grits/utils.py b/grits/utils.py index 16d4464..ca5fb56 100644 --- a/grits/utils.py +++ b/grits/utils.py @@ -23,7 +23,7 @@ def default(self, obj): elif isinstance(obj, np.ndarray): return obj.tolist() else: - return super(NumpyEncoder, self).default(obj) + return super().default(obj) def comp_from_snapshot(snapshot, indices, length_scale=1.0, mass_scale=1.0): @@ -284,7 +284,7 @@ def get_com(particle_positions, particle_masses): return center_of_mass -def get_quaternion(n1, n0=np.array([0, 0, 1])): +def get_quaternion(n1, n0=(0, 0, 1)): """Calculate rotation quaternion from axis vectors. Calculate axis and angle of rotation given @@ -295,7 +295,7 @@ def get_quaternion(n1, n0=np.array([0, 0, 1])): ---------- n1 : numpy array numpy array that is the major axis vector. - n0 : numpy array + n0 : numpy array, default (0,0,1) numpy array that is used to define the default quaternion. Defaults to the Z-axis. @@ -305,12 +305,15 @@ def get_quaternion(n1, n0=np.array([0, 0, 1])): numpy array that tells the position of the monomer in units of a quaternion. """ + n0 = np.asarray(n0) if n1 is None: # one atom in this bead -> default quaternion warnings.warn( "get_quaternion was called with None as input!\n\ Returning default orientation." ) return np.array([0, 0, 0, 1]) + else: + n1 = np.asarray(n1) V_axis = np.cross(n0, n1) theta_numerator = np.dot(n0, n1) theta_denominator = np.linalg.norm(n0) * np.linalg.norm(n1)