diff --git a/src/multimm/config.py b/src/multimm/config.py index e0905c6..54ed8ae 100644 --- a/src/multimm/config.py +++ b/src/multimm/config.py @@ -235,12 +235,16 @@ def clean_fields(cls, data: Any) -> Any: COB_EA: float = Field(default=1.0, description="Energy strength for A compartment.") COB_EB: float = Field(default=2.0, description="Energy strength for B compartment.") SCB_USE_SUBCOMPARTMENT_BLOCKS: Boolean = Field(default=False, description="Use Subcompartment Blocks.") - SCB_DISTANCE: Optional[OpenMMQuantity] = Field(default=None, description="Block copolymer equilibrium distance for chromosomal blocks.") + SCB_DISTANCE: Optional[OpenMMQuantity] = Field( + default=None, description="Block copolymer equilibrium distance for chromosomal blocks." + ) SCB_EA1: float = Field(default=1.0, description="Energy strength for A1 compartment.") SCB_EA2: float = Field(default=1.33, description="Energy strength for A2 compartment.") SCB_EB1: float = Field(default=1.66, description="Energy strength for B1 compartment.") SCB_EB2: float = Field(default=2.0, description="Energy strength for B2 compartment.") - IBL_USE_B_LAMINA_INTERACTION: Boolean = Field(default=False, description="Interactions of B compartment with lamina.") + IBL_USE_B_LAMINA_INTERACTION: Boolean = Field( + default=False, description="Interactions of B compartment with lamina." + ) IBL_SCALE: float = Field(default=400.0, description="Scaling factor for B comoartment interaction with lamina.") CF_USE_CENTRAL_FORCE: Boolean = Field(default=False, description="Attraction of smaller chromosomes.") CF_STRENGTH: float = Field(default=20.0, description="Strength of Interaction") @@ -298,7 +302,8 @@ def clean_fields(cls, data: Any) -> Any: "Options: polynomial (default, handcrafted potential), " "gaussian (soft collapse kernel), " "saturating (soft-core bounded attraction)." - )) + ), + ) CENTRAL_FORCE_TYPE: str = Field( default="harmonic", @@ -309,4 +314,5 @@ def clean_fields(cls, data: Any) -> Any: "harmonic (default, quadratic confinement around R1), " "gaussian (soft nucleolar enrichment field), " "logistic (soft-core radial partitioning with smooth boundary)." - )) \ No newline at end of file + ), + ) diff --git a/src/multimm/logger.py b/src/multimm/logger.py index 60ad8ab..d9220d8 100644 --- a/src/multimm/logger.py +++ b/src/multimm/logger.py @@ -1,15 +1,16 @@ import logging import sys + class _ColorFormatter(logging.Formatter): """Pretty colored formatter for terminal logs.""" COLORS = { - "DEBUG": "\033[1;34m", # blue - "INFO": "\033[1;32m", # green + "DEBUG": "\033[1;34m", # blue + "INFO": "\033[1;32m", # green "WARNING": "\033[1;33m", # yellow - "ERROR": "\033[1;31m", # red - "CRITICAL": "\033[1;41m", # red background + "ERROR": "\033[1;31m", # red + "CRITICAL": "\033[1;41m", # red background } RESET = "\033[0m" @@ -17,7 +18,7 @@ class _ColorFormatter(logging.Formatter): def format(self, record): level_color = self.COLORS.get(record.levelname, "") - + record.levelname = f"{level_color}{record.levelname:<8}{self.RESET}" record.name = f"\033[1;35m{record.name}{self.RESET}" record.msg = str(record.msg) @@ -26,10 +27,7 @@ def format(self, record): def setup_logger(level=logging.INFO, debug=False): - """ - Clean, colored logger for simulation pipelines. - """ - + """Clean, colored logger for simulation pipelines.""" root = logging.getLogger() # Avoid duplicate handlers @@ -39,13 +37,11 @@ def setup_logger(level=logging.INFO, debug=False): handler = logging.StreamHandler(sys.stdout) formatter = _ColorFormatter( - "\033[1;36m[%(asctime)s]\033[0m " - "%(levelname)s " - "%(name)s: %(message)s", + "\033[1;36m[%(asctime)s]\033[0m " "%(levelname)s " "%(name)s: %(message)s", datefmt="%H:%M:%S", ) handler.setFormatter(formatter) root.setLevel(logging.DEBUG if debug else level) - root.addHandler(handler) \ No newline at end of file + root.addHandler(handler) diff --git a/src/multimm/model.py b/src/multimm/model.py index 2508951..0e8b478 100644 --- a/src/multimm/model.py +++ b/src/multimm/model.py @@ -10,9 +10,8 @@ from .initial_structure_tools import build_init_mmcif, write_cmm, write_mmcif_chrom from .nucleosome_interpolation import NucleosomeInterpolation -from .utils import * from .plots import * - +from .utils import * logger = logging.getLogger(__name__) @@ -413,10 +412,7 @@ def add_chromosomal_blocks(self): logger.info("Using polynomial chromosomal self-attraction model") - self.chrom_block_force.setEnergyFunction( - "E*(k_C*r^4 - r^3 + r^2); " - "E = dE*delta(chrom1-chrom2)" - ) + self.chrom_block_force.setEnergyFunction("E*(k_C*r^4 - r^3 + r^2); " "E = dE*delta(chrom1-chrom2)") logger.info(f"k_C={self.args.CHB_KC}, dE={self.args.CHB_DE}") @@ -425,10 +421,7 @@ def add_chromosomal_blocks(self): logger.info("Using Gaussian chromosomal collapse kernel") - self.chrom_block_force.setEnergyFunction( - "-E * exp(-k_C*r^2); " - "E = dE*delta(chrom1-chrom2)" - ) + self.chrom_block_force.setEnergyFunction("-E * exp(-k_C*r^2); " "E = dE*delta(chrom1-chrom2)") logger.info(f"k_C={self.args.CHB_KC}, dE={self.args.CHB_DE}") @@ -437,10 +430,7 @@ def add_chromosomal_blocks(self): logger.info("Using saturating chromosomal interaction model") - self.chrom_block_force.setEnergyFunction( - "-E / (1 + k_C*r^2); " - "E = dE*delta(chrom1-chrom2)" - ) + self.chrom_block_force.setEnergyFunction("-E / (1 + k_C*r^2); " "E = dE*delta(chrom1-chrom2)") logger.info(f"k_C={self.args.CHB_KC}, dE={self.args.CHB_DE}") @@ -497,12 +487,10 @@ def add_Blamina_interaction(self): # 1. DEFAULT: sinusoidal shell (original) if mode == "sin": - + logger.info("Using sinusoidal lamina shell model") - self.Blamina_force.setEnergyFunction( - "B*(sin(pi*(r-R1)/(R2-R1))^8 - 1)*(delta(s+1)+delta(s+2)); " + r_expr - ) + self.Blamina_force.setEnergyFunction("B*(sin(pi*(r-R1)/(R2-R1))^8 - 1)*(delta(s+1)+delta(s+2)); " + r_expr) self.Blamina_force.addGlobalParameter("pi", np.pi) logger.info(f"Shell radii: R1={self.radius1}, R2={self.radius2}") @@ -511,8 +499,7 @@ def add_Blamina_interaction(self): logger.info("Using Gaussian lamina shell model (two-layer attraction)") self.Blamina_force.setEnergyFunction( - "-B*(exp(-(r-R1)^2/(2*sigma^2)) + exp(-(r-R2)^2/(2*sigma^2)))" - "*(delta(s+1)+delta(s+2)); " + r_expr + "-B*(exp(-(r-R1)^2/(2*sigma^2)) + exp(-(r-R2)^2/(2*sigma^2)))" "*(delta(s+1)+delta(s+2)); " + r_expr ) sigma = 0.1 * (self.radius2 - self.radius1) self.Blamina_force.addGlobalParameter("sigma", sigma) @@ -521,9 +508,7 @@ def add_Blamina_interaction(self): # 3. HARMONIC SHELL (pull to mid-shell) elif mode == "harmonic_shell": logger.info("Using harmonic lamina shell model (mid-shell attraction)") - self.Blamina_force.setEnergyFunction( - "B*(r - r0)^2*(delta(s+1)+delta(s+2)); " + r_expr - ) + self.Blamina_force.setEnergyFunction("B*(r - r0)^2*(delta(s+1)+delta(s+2)); " + r_expr) r0 = 0.5 * (self.radius1 + self.radius2) self.Blamina_force.addGlobalParameter("r0", r0) logger.info(f"r0 (mid-shell) = {r0}") @@ -532,8 +517,7 @@ def add_Blamina_interaction(self): elif mode == "logistic_shell": logger.info("Using logistic lamina shell model (soft boundaries)") self.Blamina_force.setEnergyFunction( - "-B*(1/(1+exp((r-R2)/lambda)) + 1/(1+exp(-(r-R1)/lambda)))" - "*(delta(s+1)+delta(s+2)); " + r_expr + "-B*(1/(1+exp((r-R2)/lambda)) + 1/(1+exp(-(r-R1)/lambda)))" "*(delta(s+1)+delta(s+2)); " + r_expr ) lam = 0.05 * (self.radius2 - self.radius1) self.Blamina_force.addGlobalParameter("lambda", lam) @@ -550,10 +534,7 @@ def add_Blamina_interaction(self): self.system.addForce(self.Blamina_force) def add_central_force(self): - """ - Central nucleolar attraction with chromosome-size bias. - """ - + """Central nucleolar attraction with chromosome-size bias.""" mode = getattr(self.args, "CENTRAL_FORCE_TYPE", "harmonic") self.central_force = mm.CustomExternalForce("0") @@ -581,9 +562,7 @@ def add_central_force(self): if mode == "harmonic": logger.info("Using harmonic central attraction") - self.central_force.setEnergyFunction( - f"G*chrom_s*({r}-R1)*({r}-R1)" - ) + self.central_force.setEnergyFunction(f"G*chrom_s*({r}-R1)*({r}-R1)") # ============================================================ # 2. GAUSSIAN CENTER @@ -594,9 +573,7 @@ def add_central_force(self): sigma = 0.5 * self.radius1 self.central_force.addGlobalParameter("sigma", sigma) - self.central_force.setEnergyFunction( - f"-G*chrom_s*exp(-({r}*{r})/(2*sigma*sigma))" - ) + self.central_force.setEnergyFunction(f"-G*chrom_s*exp(-({r}*{r})/(2*sigma*sigma))") # ============================================================ # 3. LOGISTIC CORE @@ -607,9 +584,7 @@ def add_central_force(self): lam = 0.2 * self.radius1 self.central_force.addGlobalParameter("lambda", lam) - self.central_force.setEnergyFunction( - f"-G*chrom_s*(1/(1+exp(({r}-R1)/lambda)))" - ) + self.central_force.setEnergyFunction(f"-G*chrom_s*(1/(1+exp(({r}-R1)/lambda)))") else: raise ValueError(f"Unknown CENTRAL_FORCE_TYPE: {mode}") @@ -636,15 +611,13 @@ def add_harmonic_bonds(self): self.system.addForce(self.bond_force) def add_loops(self): - """ - Loop constraints using stable polymer bond models. + """Loop constraints using stable polymer bond models. Supported modes: - harmonic (default) - fene_safe (bounded FENE-like) - gaussian_tether (fully smooth bounded well) """ - mode = getattr(self.args, "LE_LOOP_FORCE_TYPE", "harmonic") # 1. HARMONIC (unchanged baseline) @@ -661,9 +634,7 @@ def add_loops(self): # 2. SAFE FENE-LIKE (bounded, no singularity) elif mode == "fene_soft": - self.loop_force = mm.CustomBondForce( - "k * (r - r0)^2 / (1 + alpha * (r - r0)^2)" - ) + self.loop_force = mm.CustomBondForce("k * (r - r0)^2 / (1 + alpha * (r - r0)^2)") self.loop_force.addPerBondParameter("r0") self.loop_force.addPerBondParameter("k") @@ -675,16 +646,14 @@ def add_loops(self): r0 = self.args.LE_HARMONIC_BOND_R0 if self.args.LE_FIXED_DISTANCES else self.ds[i] k = self.args.LE_HARMONIC_BOND_K - alpha = 1.0 / (r0 ** 2) + alpha = 1.0 / (r0**2) self.loop_force.addBond(m, n, [r0, k, alpha]) # 3. GAUSSIAN TETHER (fully smooth bounded interaction) elif mode == "gaussian_tether": - self.loop_force = mm.CustomBondForce( - "k * (1 - exp(-(r - r0)^2 / sigma^2))" - ) + self.loop_force = mm.CustomBondForce("k * (1 - exp(-(r - r0)^2 / sigma^2))") self.loop_force.addPerBondParameter("r0") self.loop_force.addPerBondParameter("k") @@ -724,9 +693,7 @@ def initialize_simulation(self): logger.info("Creating initial structure...") structure_type = ( - "compartments" - if (self.Cs is not None and len(np.unique(self.Cs)) <= 3) - else "subcompartments" + "compartments" if (self.Cs is not None and len(np.unique(self.Cs)) <= 3) else "subcompartments" ) logger.info(f"Detected structure type: {structure_type}") @@ -811,7 +778,6 @@ def initialize_simulation(self): def add_forcefield(self): """Here we define the forcefield of MultiMM.""" - logger.info("Importing forcefield...") if self.args.EV_USE_EXCLUDED_VOLUME: @@ -929,11 +895,7 @@ def run_md(self): self.simulation.step(self.args.SIM_SAMPLING_STEP) - state = self.simulation.context.getState( - getPositions=True, - getEnergy=True, - getVelocities=True - ) + state = self.simulation.context.getState(getPositions=True, getEnergy=True, getVelocities=True) # STEP (always safe) step = state.getStepCount() @@ -986,10 +948,7 @@ def run_md(self): self.state.getPositions(), open(self.save_path + "model/MultiMM_afterMD.cif", "w"), ) - plot_md_thermo( - self.md_history, - self.save_path - ) + plot_md_thermo(self.md_history, self.save_path) logger.info( f"Everything is done! Simulation finished succesfully!\nMD finished in {elapsed//3600:.0f} hours, {elapsed%3600//60:.0f} minutes and {elapsed%60:.0f} seconds. ---\n" ) @@ -1079,7 +1038,8 @@ def make_plots(self): is_comp = self.Cs is not None and len(self.Cs) > 0 def _viz_and_heat(cif_path, out_name, colors=None): - """Unified structure + heatmap pipeline (single source of truth).""" + """Unified structure + heatmap pipeline (single source of + truth).""" V = get_coordinates_cif(cif_path) # 3D structure @@ -1092,14 +1052,8 @@ def _viz_and_heat(cif_path, out_name, colors=None): ) # heatmap (always) - if self.args.N_BEADS<50000: - get_heatmap( - cif_path, - viz=True, - save=True, - save_path=self.save_path + f"plots", - name=out_name - ) + if self.args.N_BEADS < 50000: + get_heatmap(cif_path, viz=True, save=True, save_path=self.save_path + f"plots", name=out_name) else: logger.warning("\033[93mHeatmap creation skipped because system is too large for visualization.\033[0m") @@ -1110,7 +1064,8 @@ def _viz_and_heat(cif_path, out_name, colors=None): name=out_name, ) - plot_projection( + if is_comp: + plot_projection( get_coordinates_mm(self.state.getPositions()), self.Cs, save_path=self.save_path, diff --git a/src/multimm/plots.py b/src/multimm/plots.py index 6fb5b9c..e47c867 100644 --- a/src/multimm/plots.py +++ b/src/multimm/plots.py @@ -1,18 +1,20 @@ import logging import os + +import matplotlib.colors as mcolors import matplotlib.pyplot as plt import numpy as np import pandas as pd import pyvista as pv import seaborn as sns +from matplotlib.lines import Line2D from matplotlib.pyplot import figure +from mpl_toolkits.mplot3d import Axes3D from scipy.spatial import ConvexHull, distance +from scipy.stats import gaussian_kde from sklearn.decomposition import PCA from sklearn.manifold import TSNE -from scipy.stats import gaussian_kde -from matplotlib.lines import Line2D -import matplotlib.colors as mcolors -from mpl_toolkits.mplot3d import Axes3D + from .utils import get_coordinates_cif logger = logging.getLogger(__name__) @@ -22,9 +24,9 @@ color_dict = {-2: "#bf0020", -1: "#e36a24", 1: "#20c8e6", 2: "#181385", 0: "#ffffff"} comp_dict = {-2: "B2", -1: "B1", 1: "A2", 2: "A1", 0: "no compartment"} + def plot_projection(struct_3D, Cs, save_path): - """ - Chromatin structural analysis centered on COM. + """Chromatin structural analysis centered on COM. Includes: - PCA embedding @@ -34,12 +36,8 @@ def plot_projection(struct_3D, Cs, save_path): - density landscapes - subcompartment-dependent structure """ - sns.set_style("whitegrid") - plt.rcParams.update({ - "figure.dpi": 600, - "font.size": 11 - }) + plt.rcParams.update({"figure.dpi": 600, "font.size": 11}) # preprocessing (STRICT ALIGNMENT GUARANTEE) X = np.asarray(struct_3D, dtype=np.float64) @@ -72,16 +70,18 @@ def plot_projection(struct_3D, Cs, save_path): eigvals = np.linalg.eigvalsh(G) anisotropy_scalar = np.sqrt(eigvals.max() / (eigvals.min() + 1e-12)) - df = pd.DataFrame({ - "x": Xc[:, 0], - "y": Xc[:, 1], - "z": Xc[:, 2], - "pc1": X_pca[:, 0], - "pc2": X_pca[:, 1], - "r_com": r, - "anisotropy": anisotropy_scalar, - "subcomp": Cs - }) + df = pd.DataFrame( + { + "x": Xc[:, 0], + "y": Xc[:, 1], + "z": Xc[:, 2], + "pc1": X_pca[:, 0], + "pc2": X_pca[:, 1], + "r_com": r, + "anisotropy": anisotropy_scalar, + "subcomp": Cs, + } + ) #df = df[df["subcomp"] != 0] @@ -111,8 +111,7 @@ def save(fig, name): fig = plt.figure(figsize=(8, 7)) ax = fig.add_subplot(111, projection="3d") - sc = ax.scatter(Xc[:, 0], Xc[:, 1], Xc[:, 2], - c=df.subcomp, cmap="Spectral", s=4, alpha=0.7) + sc = ax.scatter(Xc[:, 0], Xc[:, 1], Xc[:, 2], c=df.subcomp, cmap="Spectral", s=4, alpha=0.7) cbar = fig.colorbar(sc, ax=ax, shrink=0.6) cbar.set_label("Subcompartment state") @@ -125,8 +124,7 @@ def save(fig, name): # 3. Radial compaction (COM-based) fig, ax = plt.subplots(figsize=(7, 4)) - sns.kdeplot(data=df, x="r_com", hue="subcomp", - fill=True, alpha=0.5, palette="Spectral", ax=ax) + sns.kdeplot(data=df, x="r_com", hue="subcomp", fill=True, alpha=0.5, palette="Spectral", ax=ax) ax.set_title("Radial Compaction from Center of Mass") ax.set_xlabel("Distance from COM") @@ -136,8 +134,7 @@ def save(fig, name): # 4. PCA density landscape fig, ax = plt.subplots(figsize=(7, 6)) - sns.kdeplot(x=df.pc1, y=df.pc2, - cmap="mako", fill=True, levels=60, ax=ax) + sns.kdeplot(x=df.pc1, y=df.pc2, cmap="mako", fill=True, levels=60, ax=ax) ax.set_title("Free-energy-like landscape (PCA space)") ax.set_xlabel("PC1") @@ -151,23 +148,9 @@ def save(fig, name): norm = mcolors.Normalize(vmin=-abs_max, vmax=abs_max) cmap = plt.get_cmap("coolwarm") sns.violinplot( - data=df, - x="subcomp", - y="r_com", - palette=[cmap(norm(v)) for v in unique_sub], - inner=None, - cut=0, - ax=ax - ) - sns.stripplot( - data=df, - x="subcomp", - y="r_com", - color="black", - alpha=0.25, - size=1.5, - ax=ax + data=df, x="subcomp", y="r_com", palette=[cmap(norm(v)) for v in unique_sub], inner=None, cut=0, ax=ax ) + sns.stripplot(data=df, x="subcomp", y="r_com", color="black", alpha=0.25, size=1.5, ax=ax) ax.set_title("Radial Distribution by Subcompartment (COM frame)") ax.set_xlabel("Subcompartment state") ax.set_ylabel("Distance from COM") @@ -182,27 +165,9 @@ def save(fig, name): ] for ax, (a, b, title) in zip(axes, pairs): # replace scatter with 2D density (structure, not noise) - sns.kdeplot( - x=a, - y=b, - ax=ax, - fill=True, - cmap="mako", - levels=40, - thresh=0.05, - bw_adjust=0.6 - ) + sns.kdeplot(x=a, y=b, ax=ax, fill=True, cmap="mako", levels=40, thresh=0.05, bw_adjust=0.6) # light contour overlay (adds geometry readability) - sns.kdeplot( - x=a, - y=b, - ax=ax, - color="white", - levels=6, - linewidths=0.6, - alpha=0.6, - bw_adjust=0.6 - ) + sns.kdeplot(x=a, y=b, ax=ax, color="white", levels=6, linewidths=0.6, alpha=0.6, bw_adjust=0.6) ax.set_title(title) ax.set_xlabel("") ax.set_ylabel("") @@ -213,17 +178,14 @@ def save(fig, name): fig.suptitle("Coordinate Correlations in COM frame (density representation)", y=1.02) save(fig, "axis_correlations") - # 8. PCA KDE per subcompartment (signed colors, white background) + # 8. PCA KDE per subcompartment (signed colors, white background) fig, ax = plt.subplots(figsize=(7, 6)) # enforce clean white background (important for KDE readability) fig.patch.set_facecolor("white") ax.set_facecolor("white") x = df.pc1.values y = df.pc2.values - Xg, Yg = np.mgrid[ - x.min():x.max():200j, - y.min():y.max():200j - ] + Xg, Yg = np.mgrid[x.min() : x.max() : 200j, y.min() : y.max() : 200j] pos = np.vstack([Xg.ravel(), Yg.ravel()]) unique_sub = np.sort(df.subcomp.unique()) @@ -246,38 +208,23 @@ def save(fig, name): color = cmap(norm(scv)) - ax.contourf( - Xg, Yg, Z, - levels=3, - alpha=0.10, - colors=[color] - ) + ax.contourf(Xg, Yg, Z, levels=3, alpha=0.10, colors=[color]) - ax.contour( - Xg, Yg, Z, - levels=5, - colors=[color], - linewidths=1.2, - alpha=0.9 - ) + ax.contour(Xg, Yg, Z, levels=5, colors=[color], linewidths=1.2, alpha=0.9) # ------------------------------------------------------------ # legend (sign-based meaning preserved) # ------------------------------------------------------------ - legend_elements = [ - Line2D([0], [0], color=cmap(norm(v)), lw=2, label=f"subcomp {v}") - for v in unique_sub if v != 0 - ] + legend_elements = [Line2D([0], [0], color=cmap(norm(v)), lw=2, label=f"subcomp {v}") for v in unique_sub if v != 0] ax.legend(handles=legend_elements, frameon=True, fontsize=9) ax.set_title("Subcompartment density in PCA space") ax.set_xlabel("PC1 (collective chromatin mode)") ax.set_ylabel("PC2 (collective chromatin mode)") save(fig, "pca_kde_subcomp") + def _save_plotter(plotter, save_path): - """ - Save PyVista scene in multiple formats. - """ + """Save PyVista scene in multiple formats.""" os.makedirs(os.path.dirname(save_path), exist_ok=True) # PyVista supports direct image export @@ -299,15 +246,8 @@ def polyline_from_points(points): def viz_structure(V, colors=None, r=0.1, cmap="coolwarm", save_path=None): - """ - Visualize structure V and optionally save it to a file. - """ - - logger.info( - f"Visualizing structure: N={len(V)}, " - f"colored={colors is not None}, " - f"save_path={save_path}" - ) + """Visualize structure V and optionally save it to a file.""" + logger.info(f"Visualizing structure: N={len(V)}, " f"colored={colors is not None}, " f"save_path={save_path}") polyline = polyline_from_points(V) polyline["scalars"] = np.arange(polyline.n_points) @@ -318,6 +258,7 @@ def viz_structure(V, colors=None, r=0.1, cmap="coolwarm", save_path=None): logger.info("Color mapping enabled (signed scheme: neg/zero/pos)") + logger.info(f"Color mapping enabled: min={colors_min}, max={colors_max}, diff={diff}") # ------------------------------------------------------------ # NEW: signed piecewise normalization # ------------------------------------------------------------ @@ -386,17 +327,14 @@ def viz_structure(V, colors=None, r=0.1, cmap="coolwarm", save_path=None): logger.info("Visualization finished") + def save_chimera_cmd(start, end, total_residues, cmd_filename="coloring.cmd"): """ Create a Chimera .cmd file: - Color residues outside the given region blue. - Color residues inside the region red. """ - - logger.info( - f"Writing Chimera cmd: {cmd_filename} | " - f"region={start}-{end} | total_residues={total_residues}" - ) + logger.info(f"Writing Chimera cmd: {cmd_filename} | " f"region={start}-{end} | total_residues={total_residues}") with open(cmd_filename, "w") as f: @@ -417,6 +355,7 @@ def save_chimera_cmd(start, end, total_residues, cmd_filename="coloring.cmd"): logger.info("Chimera cmd file written successfully") + def viz_gene_structure(V, start, end, r=0.1, cmap="coolwarm", save_path=None): """Visualize structure V, highlight a continuous region in red, rest in blue.""" @@ -501,6 +440,7 @@ def viz_chroms(sim_path, r=0.1, comps=True): logger.info("Chromosome visualization finished successfully") + def get_heatmap( cif_file, viz=False, @@ -510,12 +450,9 @@ def get_heatmap( vmin=None, log_scale=True, reorder_by_diagonal=False, - name="structure" + name="structure", ): - """ - Compute and visualize contact/interaction heatmap from 3D structure. - """ - + """Compute and visualize contact/interaction heatmap from 3D structure.""" # ------------------------------------------------------------ # Output dir (UNCHANGED) # ------------------------------------------------------------ @@ -538,12 +475,9 @@ def _save_local(fig, name): # Distance → contact # ------------------------------------------------------------ mat = distance.cdist(V, V, metric="euclidean") - mat = 1.0 / (mat + 1)**(2/3) + mat = 1.0 / (mat + 1) ** (2 / 3) - logger.info( - f"Raw contact matrix: min={mat.min():.3e}, max={mat.max():.3e}, " - f"mean={mat.mean():.3e}" - ) + logger.info(f"Raw contact matrix: min={mat.min():.3e}, max={mat.max():.3e}, " f"mean={mat.mean():.3e}") if log_scale: mat = np.log1p(mat) @@ -566,12 +500,7 @@ def _save_local(fig, name): fig, ax = plt.subplots(figsize=(8, 7), dpi=600) - im = ax.imshow( - mat, - cmap="coolwarm", - interpolation="nearest", - norm=PowerNorm(gamma=0.4) - ) + im = ax.imshow(mat, cmap="coolwarm", interpolation="nearest", norm=PowerNorm(gamma=0.4)) ax.set_title("Structure-derived Contact Map", fontsize=12) ax.set_xlabel("Bead index") @@ -595,11 +524,9 @@ def _save_local(fig, name): logger.info("Heatmap computation finished") return mat -def plot_md_thermo(history, save_path): - """ - Plot energy + temperature evolution from MD. - """ +def plot_md_thermo(history, save_path): + """Plot energy + temperature evolution from MD.""" logger.info("Creating MD thermodynamics plot...") steps = history["step"] @@ -627,15 +554,14 @@ def plot_md_thermo(history, save_path): logger.info(f"MD thermodynamics plot saved to: {out}") + def analyze_structure(V, save_path, name="structure"): - """ - Advanced structural analysis for polymer-like 3D structures. + """Advanced structural analysis for polymer-like 3D structures. Outputs: - detailed report (text) - multiple physically meaningful plots """ - # ------------------------------------------------------------ # safety # ------------------------------------------------------------ @@ -693,9 +619,7 @@ def _save_local(fig, fname): v1 = V[1:-1] - V[:-2] v2 = V[2:] - V[1:-1] - cos_angles = np.sum(v1 * v2, axis=1) / ( - np.linalg.norm(v1, axis=1) * np.linalg.norm(v2, axis=1) + 1e-8 - ) + cos_angles = np.sum(v1 * v2, axis=1) / (np.linalg.norm(v1, axis=1) * np.linalg.norm(v2, axis=1) + 1e-8) angles = np.arccos(np.clip(cos_angles, -1, 1)) @@ -718,9 +642,9 @@ def _save_local(fig, fname): local_rg = [] for i in range(N - window): - chunk = V[i:i + window] + chunk = V[i : i + window] cm = np.mean(chunk, axis=0) - local_rg.append(np.sqrt(np.mean(np.sum((chunk - cm)**2, axis=1)))) + local_rg.append(np.sqrt(np.mean(np.sum((chunk - cm) ** 2, axis=1)))) local_rg = np.array(local_rg) @@ -759,14 +683,14 @@ def _save_local(fig, fname): f.write("Local Rg → domain compaction\n") # PLOTS - os.makedirs(base+'/plots', exist_ok=True) + os.makedirs(base + "/plots", exist_ok=True) # 1. bond lengths fig, ax = plt.subplots(figsize=(6, 4)) ax.hist(bonds, bins=80) ax.set_title("Bond Length Distribution") ax.set_xlabel("Bond length") - _save_local(fig, base+f"/plots/{name}_bonds") + _save_local(fig, base + f"/plots/{name}_bonds") # ------------------------------------------------------------ @@ -775,7 +699,7 @@ def _save_local(fig, fname): ax.hist(angles, bins=80) ax.set_title("Angle Distribution") ax.set_xlabel("Angle (rad)") - _save_local(fig, base+f"/plots/{name}_angles") + _save_local(fig, base + f"/plots/{name}_angles") # ------------------------------------------------------------ @@ -785,7 +709,7 @@ def _save_local(fig, fname): ax.set_title("Distance vs Genomic Separation") ax.set_xlabel("Genomic separation (beads)") ax.set_ylabel("Mean spatial distance") - _save_local(fig, base+f"/plots/{name}_scaling") + _save_local(fig, base + f"/plots/{name}_scaling") # ------------------------------------------------------------ @@ -795,7 +719,7 @@ def _save_local(fig, fname): ax.set_title("Scaling (log-log)") ax.set_xlabel("s") ax.set_ylabel("R(s)") - _save_local(fig, base+f"/plots/{name}_scaling_loglog") + _save_local(fig, base + f"/plots/{name}_scaling_loglog") # ------------------------------------------------------------ @@ -805,7 +729,7 @@ def _save_local(fig, fname): ax.set_title("Local Compaction (Sliding Rg)") ax.set_xlabel("Bead index") ax.set_ylabel("Local Rg") - _save_local(fig, base+f"/plots/{name}_local_compaction") + _save_local(fig, base + f"/plots/{name}_local_compaction") # ------------------------------------------------------------ @@ -816,7 +740,7 @@ def _save_local(fig, fname): ax.hist(r, bins=60) ax.set_title("Radial Distribution") ax.set_xlabel("Distance from COM") - _save_local(fig, base+f"/plots/{name}_radial") + _save_local(fig, base + f"/plots/{name}_radial") # ------------------------------------------------------------ return { @@ -826,4 +750,4 @@ def _save_local(fig, fname): "density": density, "asphericity": asphericity, "acylindricity": acylindricity, - } \ No newline at end of file + } diff --git a/src/multimm/run.py b/src/multimm/run.py index 7144d68..316e91f 100644 --- a/src/multimm/run.py +++ b/src/multimm/run.py @@ -13,9 +13,9 @@ from openmm.unit import Quantity from .config import SimulationConfig +from .logger import setup_logger from .model import MultiMM from .utils import chrom_sizes -from .logger import setup_logger setup_logger() logger = logging.getLogger(__name__) @@ -49,14 +49,12 @@ def print_startup_banner(logger): - """ - Print colorful MultiMM startup banner. - """ - + """Print colorful MultiMM startup banner.""" for i, line in enumerate(STARTUP_BANNER_LINES): color = COLORS[i % len(COLORS)] logger.info(f"{color}{line}{RESET}") + class Tee: def __init__(self, *streams): self.streams = streams @@ -70,6 +68,7 @@ def flush(self): for s in self.streams: s.flush() + class ArgumentChanger: # ------------------------------------------------------------ @@ -87,9 +86,7 @@ def __init__(self, args, chrom_sizes): self._original_values = {} def set_arg(self, name, value): - """ - Set argument value in attribute and store change history. - """ + """Set argument value in attribute and store change history.""" if hasattr(self.args, name): old_value = getattr(self.args, name, None) @@ -103,9 +100,7 @@ def set_arg(self, name, value): logger.warning(f"Argument '{name}' not found in args object.") def _report_changes(self): - """ - Print all modified parameters in a readable way. - """ + """Print all modified parameters in a readable way.""" if not self._original_values: return @@ -156,9 +151,7 @@ def convenient_argument_changer(self): elif level in ("region", "loc"): logger.warning( - f"{self.ORANGE}{self.BOLD}" - "Region-level modelling activated. Overwriting parameters." - f"{self.RESET}" + f"{self.ORANGE}{self.BOLD}" "Region-level modelling activated. Overwriting parameters." f"{self.RESET}" ) self.set_arg("N_BEADS", 5000) @@ -194,9 +187,7 @@ def convenient_argument_changer(self): elif level in ("gw", "genome"): logger.warning( - f"{self.ORANGE}{self.BOLD}" - "Genome-wide modelling activated. Overwriting parameters." - f"{self.RESET}" + f"{self.ORANGE}{self.BOLD}" "Genome-wide modelling activated. Overwriting parameters." f"{self.RESET}" ) self.set_arg("N_BEADS", 200000) @@ -216,11 +207,12 @@ def convenient_argument_changer(self): if self.args.MODELLING_LEVEL: self._report_changes() + def args_tests(args): def check_file(path, name, ext_hint=None): - """ - Validate optional input files. + """Validate optional input files. + If provided, ensure they exist. """ if path is None or path == "": @@ -228,9 +220,7 @@ def check_file(path, name, ext_hint=None): if not os.path.exists(path): ext_msg = f" (expected {ext_hint})" if ext_hint else "" - raise ValueError( - f"\033[91m{name} file was provided but not found: {path}{ext_msg}\033[0m" - ) + raise ValueError(f"\033[91m{name} file was provided but not found: {path}{ext_msg}\033[0m") # ----------------------------------------- # REQUIRED INPUT @@ -251,8 +241,7 @@ def check_file(path, name, ext_hint=None): if args.LOOPS_PATH is None or args.LOOPS_PATH == "": raise ValueError( - "\033[91mInteraction data is required to run MultiMM." - "Please provide a .bedpe file via LOOPS_PATH.\033[0m" + "\033[91mInteraction data is required to run MultiMM." "Please provide a .bedpe file via LOOPS_PATH.\033[0m" ) elif (args.COMPARTMENT_PATH is None or args.COMPARTMENT_PATH == "") and args.COB_USE_COMPARTMENT_BLOCKS: diff --git a/src/multimm/utils.py b/src/multimm/utils.py index 86fe26f..ab8795b 100644 --- a/src/multimm/utils.py +++ b/src/multimm/utils.py @@ -167,7 +167,6 @@ def get_coordinates_mm(mm_vec): def get_coordinates_cif(file): """It returns the coordinate matrix V (N,3) from a .cif/.pdb-like file.""" - logger.info(f"Loading structure file: {file}") V = [] @@ -197,13 +196,11 @@ def get_coordinates_cif(file): V = np.array(V) - logger.info( - f"Structure loaded: atoms={n_atoms}, " - f"total_lines={n_lines}, shape={V.shape}" - ) + logger.info(f"Structure loaded: atoms={n_atoms}, " f"total_lines={n_lines}, shape={V.shape}") return V + def compute_averages(arr1, N2): # Calculate the window size window_size = len(arr1) // N2 @@ -393,6 +390,7 @@ def write_chrom_colors( logger.info("Chromosome color file written successfully") + def min_max_trans(x): return (x - x.min()) / (x.max() - x.min()) @@ -528,7 +526,7 @@ def import_mns_from_bedpe( if down_prob < 1.0: ms, ns, cs, ds = downsample_arrays(ms, ns, cs, ds, down_prob) - + avg_ls = np.average(ns - ms) logger.info(f"Average loop size: {avg_ls}") diff --git a/tests/test_simulations.py b/tests/test_simulations.py index aedfc13..4b26314 100644 --- a/tests/test_simulations.py +++ b/tests/test_simulations.py @@ -127,3 +127,24 @@ def test_subprocess_run(tmp_path): assert os.path.exists(os.path.join(str(out_dir), "metadata", "output.log")) +def test_simulation_region_no_compartments_with_plots(tmp_path): + out_dir = tmp_path / "sim_region_no_comps" + config = SimulationConfig( + LOOPS_PATH="tests/fixtures/ENCFF045MJY_simple.bedpe", + COMPARTMENT_PATH=None, + OUT_PATH=str(out_dir), + N_BEADS=100, + CHROM="chr1", + LOC_START=16000000, + LOC_END=17950000, + SIM_RUN_MD=True, + SIM_N_STEPS=5, + SAVE_PLOTS=True, + COB_USE_COMPARTMENT_BLOCKS=False, + ) + md = MultiMM(config) + md.run() + assert os.path.exists(os.path.join(str(out_dir), "model", "MultiMM_minimized.cif")) + + +