EntDetect 1.2.0__py3-none-any.whl

scripts/run_change_resolution.py

@@ -0,0 +1,63 @@
+"""
+python scripts/run_change_resolution.py 
+--outdir TestingGrounds/CoarseGraining/ 
+--pdbfile /path/to/structure.pdb 
+--nscal 2 
+--domain_file /path/to/domain_def.dat 
+--ID test_protein
+"""
+
+def main(argv=None):
+
+    import sys, os
+    import argparse
+    import time
+
+    start_time = time.time()
+
+    parser = argparse.ArgumentParser(description="Process user specified arguments")
+    parser.add_argument("--outdir", type=str, required=True, help="output directory for results")
+    parser.add_argument("--pdbfile", type=str, required=True, help="Path to the all-atom PDB file")
+    parser.add_argument("--nscal", type=int, default=2, help="Coarse graining scale factor")
+    parser.add_argument("--domain_file", type=str, required=True, help="Path to the domain definition file")
+    parser.add_argument("--ID", type=str, required=True, help="Tag for the coarse graining process")
+    args = parser.parse_args(argv)
+    print(args)
+    outdir = args.outdir
+    pdbfile = args.pdbfile
+    nscal = args.nscal
+    domain_file = args.domain_file
+    ID = args.ID
+
+    # Import here so `run_change_resolution -h` works without OpenMM installed.
+    from EntDetect.change_resolution import CoarseGrain, BackMapping
+
+    ## Coarse grain the all-atom structure
+    CoarseGrainer = CoarseGrain(outdir=outdir,
+                                 ID=ID,
+                                 pdbfile=pdbfile,
+                                 nscal=nscal, 
+                                 domain_file=domain_file)
+    print(CoarseGrainer)
+   
+    # run the coarse graining for the all-atom pdbfile
+    CGfiles = CoarseGrainer.run()
+    print(CGfiles)
+
+    # parse the prm and top file to make a OpenMM compatible .xml force field file
+    CoarseGrainer.parse_cg_prm(prmfile=CGfiles['prm'], topfile=CGfiles['top'])
+
+
+    ## Backmap the coarse grained structure to all-atom
+    backMapper = BackMapping(outdir=outdir)
+    print(f'BackMapper: {backMapper}')
+
+    # backmap (using original pdbfile as aa_pdb reference)
+    backMapper.backmap(cg_pdb=CGfiles['cor'], aa_pdb=pdbfile, TAG=ID)
+
+    print(f'NORMAL TERMINATION - {time.time() - start_time} seconds')
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

scripts/run_compare_sim2exp.py

@@ -0,0 +1,215 @@
+from EntDetect.gaussian_entanglement import GaussianEntanglement
+from EntDetect.clustering import ClusterNativeEntanglements, MSMNonNativeEntanglementClustering
+from EntDetect.order_params import CalculateOP, CollectOP
+from EntDetect.compare_sim2exp import MassSpec
+
+"""
+Collect per-trajectory SASA/XP outputs (optional) and run the LiP-MS / XL-MS
+consistency test.
+
+Two usage modes:
+
+  1. Collect + run  (provide --sasa_dir and --xp_dir):
+     python scripts/run_compare_sim2exp.py
+         --sasa_dir   /path/to/OP_AA/SASA
+         --xp_dir     /path/to/OP_AA/XP
+         --n_traj     1000
+         --n_frames   335
+         --msm_data_file ...  (remaining args as below)
+
+  2. Skip collection (provide --sasa_data_file and --dist_data_file directly):
+     python scripts/run_compare_sim2exp.py
+         --sasa_data_file /path/to/SASA.npy
+         --dist_data_file /path/to/Jwalk.npy
+         --msm_data_file ...
+
+Full example (mode 1):
+python scripts/run_compare_sim2exp.py
+--sasa_dir        /path/to/OP_AA/SASA
+--xp_dir          /path/to/OP_AA/XP
+--n_traj          1000
+--n_frames        335
+--msm_data_file   /path/to/msm_data.csv
+--meta_dist_file  /path/to/meta_dist.npy
+--LiPMS_exp_file  /path/to/LiPMS_exp.xlsx
+--XLMS_exp_file   /path/to/XLMS_exp.xlsx
+--cluster_data_file /path/to/cluster_data.npz
+--OPpath          /path/to/OP_AA/
+--AAdcd_dir       /path/to/aa_trajectories/
+--native_AA_pdb   /path/to/native.pdb
+--state_idx_list  4 6 8
+--prot_len        387
+--last_num_frames 335
+--rm_traj_list    65 75 155
+--native_state_idx 9
+--outdir          /path/to/outdir/
+--ID              1ZMR
+--start           6600
+--end             -1
+--stride          1
+--num_perm        1000
+--n_boot          100
+--lag_frame       20
+--nproc           10
+"""
+
+def main(argv=None):
+
+    import multiprocessing as mp
+    import sys, os
+    import argparse
+    import time
+
+    start_time = time.time()
+    
+    parser = argparse.ArgumentParser(description="Process user specified arguments")
+    parser.add_argument("--msm_data_file", type=str, required=True, help="Path to MSM mapping file")
+    parser.add_argument("--meta_dist_file", type=str, required=True, help="Path to meta-distance file")
+    parser.add_argument("--LiPMS_exp_file", type=str, required=True, help="Path to LiP-MS experimental data file")
+    parser.add_argument("--XLMS_exp_file", type=str, required=True, help="Path to XL-MS experimental data file")
+    parser.add_argument("--cluster_data_file", type=str, required=True, help="Path to clustering data file")
+    parser.add_argument("--OPpath", type=str, required=True, help="Path to order parameters directory")
+    parser.add_argument("--AAdcd_dir", type=str, required=True, help="Path to all-atom DCD files directory")
+    parser.add_argument("--native_AA_pdb", type=str, required=True, help="Path to native all-atom PDB file")
+    parser.add_argument("--state_idx_list", type=int, nargs='+', required=True, help="List of state indices to analyze")
+    parser.add_argument("--prot_len", type=int, required=True, help="Length of the protein")
+    parser.add_argument("--last_num_frames", type=int, required=True, help="Number of last frames to consider")
+    parser.add_argument("--rm_traj_list", type=int, nargs='+', required=True, help="List of trajectory indices to remove")
+    parser.add_argument("--native_state_idx", type=int, required=True, help="Index of the native state")
+    parser.add_argument("--outdir", type=str, required=True, help="Output directory for results")
+    parser.add_argument("--ID", type=str, required=True, help="An ID for the analysis")
+    parser.add_argument("--start", type=int, required=True, help="Start frame index")
+    parser.add_argument("--end", type=int, required=True, help="End frame index")
+    parser.add_argument("--stride", type=int, required=True, help="Stride for frame selection")
+    parser.add_argument("--verbose", action='store_true', help="Enable verbose output")
+    parser.add_argument("--num_perm", type=int, required=True, help="Number of permutations for statistical tests")
+    parser.add_argument("--n_boot", type=int, required=True, help="Number of bootstrap samples")
+    parser.add_argument("--lag_frame", type=int, required=True, help="Lag time in frames")
+    parser.add_argument("--nproc", type=int, required=True, help="Number of processes for parallel computation")
+    # --- CollectOP arguments (optional: collect from per-traj files) ---
+    parser.add_argument("--sasa_dir", type=str, default=None,
+                        help="Directory of per-traj {ID}_Traj{N}.SASA files. "
+                             "If provided together with --xp_dir, CollectOP is run "
+                             "before MassSpec and --sasa_data_file / --dist_data_file "
+                             "are set automatically.")
+    parser.add_argument("--xp_dir", type=str, default=None,
+                        help="Directory of per-traj {ID}_Traj{N}.XP files (used with --sasa_dir).")
+    parser.add_argument("--n_traj", type=int, default=None,
+                        help="Total number of trajectories for CollectOP (required with --sasa_dir).")
+    parser.add_argument("--n_frames", type=int, default=None,
+                        help="Frames per trajectory stored in each file (required with --sasa_dir).")
+    parser.add_argument("--collect_jwalk_npy", action='store_true',
+                        help="Also build Jwalk.npy with CollectOP (legacy path). "
+                             "By default, XL-MS scoring streams directly from XP files to reduce memory use.")
+    # --- Direct array paths (used when skipping collection) ---
+    parser.add_argument("--sasa_data_file", type=str, default=None,
+                        help="Path to pre-built SASA.npy. Required if --sasa_dir is not provided.")
+    parser.add_argument("--dist_data_file", type=str, default=None,
+                        help="Path to pre-built Jwalk.npy. Required if --xp_dir is not provided.")
+    args = parser.parse_args(argv)
+    print(args)
+    msm_data_file = args.msm_data_file
+    meta_dist_file = args.meta_dist_file
+    LiPMS_exp_file = args.LiPMS_exp_file
+    sasa_data_file = args.sasa_data_file
+    XLMS_exp_file = args.XLMS_exp_file
+    dist_data_file = args.dist_data_file
+    cluster_data_file = args.cluster_data_file
+    OPpath = args.OPpath
+    AAdcd_dir = args.AAdcd_dir
+    native_AA_pdb = args.native_AA_pdb
+    state_idx_list = args.state_idx_list
+    prot_len = args.prot_len
+    last_num_frames = args.last_num_frames
+    rm_traj_list = args.rm_traj_list
+    native_state_idx = args.native_state_idx
+    outdir = args.outdir
+    ID = args.ID
+    start = args.start
+    end = args.end
+    stride = args.stride
+    verbose = args.verbose
+    num_perm = args.num_perm
+    n_boot = args.n_boot
+    lag_frame = args.lag_frame
+    nproc = args.nproc
+
+    # ── validate input mode ────────────────────────────────────────────────
+    collect_mode = args.sasa_dir is not None and args.xp_dir is not None
+    direct_mode  = args.sasa_data_file is not None and args.dist_data_file is not None
+
+    if not collect_mode and not direct_mode:
+        parser.error(
+            "Provide either (--sasa_dir + --xp_dir + --n_traj + --n_frames) "
+            "to collect from per-trajectory files, or "
+            "(--sasa_data_file + --dist_data_file) to use pre-built arrays."
+        )
+
+    # ── Step 1: collect per-trajectory outputs if requested ───────────────
+    if collect_mode:
+        if args.n_traj is None or args.n_frames is None:
+            parser.error("--n_traj and --n_frames are required when using --sasa_dir / --xp_dir")
+
+        os.makedirs(outdir, exist_ok=True)
+        collector = CollectOP(
+            sasa_dir  = args.sasa_dir,
+            xp_dir    = args.xp_dir,
+            outdir    = outdir,
+            ID        = ID,
+            n_traj    = args.n_traj,
+            n_frames  = args.n_frames,
+            prot_len  = prot_len,
+        )
+        sasa_data_file = collector.collect_SASA()
+        if args.collect_jwalk_npy:
+            dist_data_file = collector.collect_Jwalk()
+        else:
+            dist_data_file = None
+        print(f'CollectOP SASA:  {sasa_data_file}')
+        if dist_data_file is not None:
+            print(f'CollectOP Jwalk: {dist_data_file}')
+        else:
+            print('CollectOP Jwalk: skipped (streaming XP mode enabled)')
+
+    # ── Step 2: run the consistency test ──────────────────────────────────
+    MS = MassSpec(msm_data_file=msm_data_file,
+                    meta_dist_file=meta_dist_file,
+                    LiPMS_exp_file=LiPMS_exp_file,
+                    sasa_data_file=sasa_data_file,
+                    XLMS_exp_file=XLMS_exp_file,
+                    dist_data_file=dist_data_file,
+                    cluster_data_file=cluster_data_file,
+                    OPpath=OPpath,
+                    AAdcd_dir=AAdcd_dir,
+                    native_AA_pdb=native_AA_pdb,
+                    xp_dir=args.xp_dir,
+                    state_idx_list=state_idx_list,
+                    prot_len=prot_len,
+                    last_num_frames=last_num_frames,
+                    rm_traj_list=rm_traj_list,
+                    native_state_idx=native_state_idx,
+                    outdir=outdir,
+                    ID=ID,
+                    start=start,
+                    end=end,
+                    stride=stride,
+                    verbose=verbose,
+                    num_perm=num_perm,
+                    n_boot=n_boot,
+                    lag_frame=lag_frame,
+                    nproc=nproc)
+
+    # run the consistency test
+    consist_data_file, consist_result_file = MS.LiP_XL_MS_ConsistencyTest()
+    print(f'consist_data_file: {consist_data_file}')
+    print(f'consist_result_file: {consist_result_file}')
+
+    # select the representative structures from the consistency test
+    MS.select_rep_structs(consist_data_file, consist_result_file, total_traj_num_frames=335, last_num_frames=67)
+    
+    print(f'NORMAL TERMINATION - {time.time() - start_time} seconds')
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

scripts/run_montecarlo.py

@@ -0,0 +1,158 @@
+#!/usr/bin/env python3
+from EntDetect.statistics import MonteCarlo
+from EntDetect._logging import setup_logger
+
+"""
+Run Workflow 4 Monte Carlo subpopulation selection.
+
+This script wraps EntDetect.statistics.MonteCarlo and optimizes population
+partitions using a logistic-regression objective and penalty terms.
+
+Example
+-------
+python scripts/run_montecarlo.py \
+    --dataframe_files /path/to/residue_dataframes_workflow4.csv \
+    --outpath /path/to/workflow4/monte_carlo/ \
+    --gene_list /path/to/gene_list.txt \
+    --tag Ecoli_population_mc \
+    --steps 100000 \
+    --n_groups 4 \
+    --C1 1.0 \
+    --C2 2.5 \
+    --beta 0.05
+
+Expected input schema
+---------------------
+- Either a single combined design matrix file OR one residue table per protein
+- Files are pipe-delimited ("|")
+- Required columns: gene, mapped_resid, uniprot_length, AA, region, cut_C_Rall
+"""
+
+def main(argv=None):
+
+    import os
+    import argparse
+    import time
+    import logging
+
+    start_time = time.time()
+
+    parser = argparse.ArgumentParser(
+        description="Run Workflow 4 Monte Carlo subpopulation selection."
+    )
+
+    # --- required IO ---
+    parser.add_argument("--dataframe_files", type=str, required=True,
+                        help="Input design matrix path: either a directory of per-protein files or a single combined CSV")
+    parser.add_argument("--outpath", type=str, required=True,
+                        help="Output directory for Monte Carlo results")
+    parser.add_argument("--gene_list", type=str, required=True,
+                        help="Path to gene list file (one ID per line)")
+    parser.add_argument("--tag", type=str, required=True,
+                        help="Identifier tag for output naming")
+
+    # --- model options ---
+    parser.add_argument("--reg_formula", type=str, default='cut_C_Rall ~ region + AA',
+                        help="Regression formula used by state scoring")
+    parser.add_argument("--response_var", type=str, default='cut_C_Rall',
+                        help="Response variable in regression")
+    parser.add_argument("--test_var", type=str, default='region',
+                        help="Primary test variable")
+    parser.add_argument("--random", action='store_true',
+                        help="Use random sampling mode")
+    parser.add_argument("--n_groups", type=int, default=4,
+                        help="Number of groups (default: 4)")
+    parser.add_argument("--steps", type=int, default=100000,
+                        help="Number of Monte Carlo steps (default: 100000)")
+    parser.add_argument("--C1", type=float, default=1.0,
+                        help="Monte Carlo objective weight C1")
+    parser.add_argument("--C2", type=float, default=2.5,
+                        help="Monte Carlo objective weight C2")
+    parser.add_argument("--beta", type=float, default=0.05,
+                        help="Inverse temperature/annealing parameter beta")
+    parser.add_argument("--linearT", action='store_true',
+                        help="Use linear temperature schedule")
+
+    # --- logging ---
+    parser.add_argument("--log_level", default="INFO", choices=["DEBUG", "INFO", "WARNING", "ERROR"],
+                        help="Logging verbosity (default: INFO)")
+    parser.add_argument("--logdir", type=str, default=None,
+                        help="Directory for log files (default: same as --outpath)")
+
+    args = parser.parse_args(argv)
+
+    dataframe_files = args.dataframe_files
+    outdir = args.outpath
+    gene_list = args.gene_list
+    tag = args.tag
+    reg_formula = args.reg_formula
+    response_var = args.response_var
+    test_var = args.test_var
+    random = args.random
+    n_groups = args.n_groups
+    steps = args.steps
+    C1 = args.C1
+    C2 = args.C2
+    beta = args.beta
+    linearT = args.linearT
+
+    log_level = getattr(logging, args.log_level.upper(), logging.INFO)
+    logdir = args.logdir if args.logdir is not None else outdir
+    os.makedirs(logdir, exist_ok=True)
+
+    logger = setup_logger('run_montecarlo', outdir=logdir, ID=tag, log_level=log_level)
+    setup_logger('MonteCarlo', outdir=logdir, ID=tag, log_level=log_level)
+    logger.info(f'args: {args}')
+
+    # --- input validation ---
+    if not os.path.exists(dataframe_files):
+        parser.error(f"--dataframe_files does not exist: {dataframe_files}")
+    if not os.path.isfile(gene_list):
+        parser.error(f"--gene_list does not exist or is not a file: {gene_list}")
+    if n_groups < 1:
+        parser.error("--n_groups must be >= 1")
+    if steps < 1:
+        parser.error("--steps must be >= 1")
+    os.makedirs(outdir, exist_ok=True)
+
+    # --- step 1: initialize Monte Carlo object ---
+    MC = MonteCarlo(
+        dataframe_files=dataframe_files,
+        outdir=outdir,
+        gene_list=gene_list,
+        ID=tag,
+        reg_formula=reg_formula,
+        response_var=response_var,
+        test_var=test_var,
+        random=random,
+        n_groups=n_groups,
+        steps=steps,
+        C1=C1,
+        C2=C2,
+        beta=beta,
+        linearT=linearT,
+        log_level=log_level,
+        logdir=logdir,
+    )
+    logger.info(f'MonteCarlo: {MC}')
+
+    # --- step 2: load residue-level data ---
+    MC.load_data(
+        sep='|',
+        reg_var=['AA', 'region'],
+        response_var='cut_C_Rall',
+        var2binarize=['cut_C_Rall', 'region'],
+        mask_column='mapped_resid',
+        ID_column='gene',
+        Length_column='uniprot_length',
+    )
+
+    # --- step 3: run simulation ---
+    MC.run(encoded_df=MC.data, ID_column='gene')
+    
+    logger.info(f'NORMAL TERMINATION - {time.time() - start_time:.1f} seconds')
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

scripts/run_nativeNCLE.py

@@ -0,0 +1,179 @@
+#!/usr/bin/env python3
+from EntDetect.gaussian_entanglement import GaussianEntanglement
+from EntDetect.clustering import ClusterNativeEntanglements
+from EntDetect.entanglement_features import FeatureGen
+from EntDetect._logging import setup_logger
+
+"""
+Script to calculate native Gaussian entanglements in a given structure (PDB or COR file),
+filter for high-quality entanglements, cluster them, and generate entanglement features.
+
+Usage example (1ZMR / ecPGK):
+    python scripts/run_nativeNCLE.py \\
+        --struct  /scratch/ims86/EntDetect_Datastore/user_input/reference_structures/1zmr_model_clean.pdb \\
+        --outdir  /scratch/ims86/EntDetect_Datastore/outputs/workflow1 \\
+        --ID      1ZMR \\
+        --chain   A \\
+        --organism Ecoli \\
+        --Accession P00558 \\
+        --model   EXP
+
+Arguments:
+    --struct              Path to input PDB (or COR) structure file                    [required]
+    --outdir              Root output directory; sub-dirs are created automatically    [required]
+    --ID                  Identifier for the analysis (default: structure basename)
+    --chain               Chain ID to process; omit to process all chains
+    --organism            Reference proteome for clustering: Ecoli | Human | Yeast     (default: Ecoli)
+    --Accession           UniProt accession used in feature-file naming                (default: P00558)
+    --model               Structure type for HQ filtering: EXP | AF                   (default: EXP)
+    --resolution          Structure resolution: aa | cg  (overrides --cg flag)
+    --contacts            Contact definition: heavy | calpha
+    --cg                  Flag: input is a coarse-grained C-alpha model (legacy; prefer --resolution cg)
+    --Calpha              Flag: use C-alpha contacts (legacy; prefer --contacts calpha)
+    --cluster_cutoff      Clustering distance cutoff in Å; if omitted, uses the
+                          organism-specific default (Ecoli: 57, Human: 52, Yeast: 49)
+    --ent_detection_method
+                          Entanglement detection criterion:
+                            1 = any nonzero GLN for either termini
+                            2 = any nonzero TLN for either termini  (class default)
+                            3 = both GLN and TLN nonzero for same termini  (recommended; script default)
+"""
+
+
+def main(argv=None):
+
+    import multiprocessing as mp
+    import sys, os
+    import argparse
+    import time
+
+    start_time = time.time()
+
+    parser = argparse.ArgumentParser(description="Process user specified arguments")
+    parser.add_argument("--struct", type=str, required=True, help="Path to PDB structure file")
+    parser.add_argument("--outdir", type=str, required=True, help="output directory for results")
+    parser.add_argument("--ID", type=str, required=False, help="An id for the analysis (defaults to structure basename)")
+    parser.add_argument("--chain", type=str, required=False, help="Chain identifier (optional, processes all chains if not specified)", default=None)
+    parser.add_argument("--organism", type=str, required=False, help="Organism name for clustering: {Ecoli, Human, Yeast}", default='Ecoli')
+    parser.add_argument("--Accession", type=str, required=False, help="UniProt Accession for the protein", default='P00558')
+    parser.add_argument("--cg", action='store_true', help="Indicate structure is coarse-grained (C-alpha only) model")
+    parser.add_argument(
+        "--Calpha",
+        "--calpha",
+        action='store_true',
+        dest="Calpha",
+        help="Use C-alpha atoms for contact definition (legacy flag; prefer --contacts calpha)",
+    )
+    parser.add_argument(
+        "--resolution",
+        type=str,
+        choices=["aa", "cg"],
+        default=None,
+        help="Structure resolution: 'aa' (all-atom) or 'cg' (C-alpha coarse-grained). If set, this overrides --cg.",
+    )
+    parser.add_argument(
+        "--contacts",
+        type=str,
+        choices=["heavy", "calpha"],
+        default=None,
+        help="Contact definition to use: 'heavy' (all-atom) or 'calpha'. If omitted, defaults to heavy for aa and calpha for cg.",
+    )
+    parser.add_argument("--cluster_cutoff", type=float, required=False,
+                        help="Clustering distance cutoff in Å. If omitted, uses the organism-specific default (Ecoli: 57, Human: 52, Yeast: 49).",
+                        default=None)
+    parser.add_argument("--model", type=str, required=False, help="Model type for high-quality selection: {EXP, AF}", default='EXP')
+    parser.add_argument("--ent_detection_method", type=int, required=False, help="ENT detection method: 1=any GLN, 2=any TLN (default), 3=both GLN and TLN same termini", default=3)
+    parser.add_argument("--log_level", type=str, default="INFO", choices=["DEBUG", "INFO", "WARNING", "ERROR"],
+                        help="Logging verbosity level (default: INFO)")
+    parser.add_argument("--logdir", type=str, default=None,
+                        help="Directory for log file. Defaults to --outdir if not specified.")
+    args = parser.parse_args(argv)
+    import logging
+    log_level = getattr(logging, args.log_level.upper(), logging.INFO)
+    
+    struct = args.struct
+    outdir = args.outdir
+    os.makedirs(outdir, exist_ok=True)
+    ID = args.ID if args.ID is not None else os.path.splitext(os.path.basename(struct))[0]
+    logdir = args.logdir if args.logdir is not None else outdir
+    # Pre-configure all EntDetect loggers for this run so they share one log file
+    logger = setup_logger('run_nativeNCLE', outdir=logdir, ID=ID, log_level=log_level)
+    for _cls in ['GaussianEntanglement', 'ClusterNativeEntanglements', 'FeatureGen']:
+        setup_logger(_cls, outdir=logdir, ID=ID, log_level=log_level)
+    logger.info(f'args: {args}')
+    chain = args.chain
+    organism = args.organism
+    cluster_cutoff = args.cluster_cutoff
+    model = args.model
+
+    # Derive effective resolution/contact settings while keeping legacy flags working.
+    # If neither --resolution nor --contacts are provided, behavior matches historical defaults:
+    #   - all-atom (CG=False)
+    #   - heavy-atom contacts (Calpha=False)
+    if args.resolution is None:
+        CG = bool(args.cg)
+    else:
+        CG = args.resolution == "cg"
+
+    if args.contacts is None:
+        Calpha = bool(args.Calpha) if args.resolution is None else (CG is True)
+    else:
+        Calpha = args.contacts == "calpha"
+
+    # Set up Gaussian Entanglement and Clustering objects
+    ge = GaussianEntanglement(g_threshold=0.6, density=1.0, Calpha=Calpha, CG=CG, ent_detection_method=args.ent_detection_method, log_level=log_level, logdir=logdir)
+    clustering = ClusterNativeEntanglements(organism=organism, cut_off=cluster_cutoff, log_level=log_level, logdir=logdir)
+
+    # Determine which chains to process
+    if chain is not None:
+        chains_to_process = [chain]
+    else:
+        # Get all chains from the structure
+        import MDAnalysis as mda
+        u = mda.Universe(struct)
+        chains_to_process = sorted(set([atom.segid if atom.segid else 'A' for atom in u.atoms if atom.segid or atom.chainID]))
+        if not chains_to_process or chains_to_process == ['']:
+            # Fallback: use mdtraj to get chains
+            import mdtraj as md
+            traj = md.load(struct)
+            chains_to_process = sorted(set([c.chain_id for c in traj.topology.chains]))
+        logger.info(f'Processing chains: {chains_to_process}')
+
+    # Process each chain separately for all steps
+    for chain_id in chains_to_process:
+        logger.info(f"{'='*80}\nProcessing chain {chain_id}\n{'='*80}")
+        
+        # Use chain suffix for file naming when processing multiple chains
+        if len(chains_to_process) > 1:
+            hq_id = f"{ID}_{chain_id}"
+        else:
+            hq_id = ID
+        
+        # All chains use the same Native_GE directory
+        ge_outdir = os.path.join(outdir, 'Native_GE')
+        os.makedirs(ge_outdir, exist_ok=True)
+        
+        # Calculate native entanglements for this chain
+        NativeEnt = ge.calculate_native_entanglements(struct, outdir=ge_outdir, ID=hq_id, chain=chain_id)
+        logger.info(f'Native entanglements saved to {NativeEnt["outfile"]}')
+        
+        # Optional steps: select high-quality entanglements 
+        HQNativeEnt = ge.select_high_quality_entanglements(NativeEnt['outfile'], struct, outdir=os.path.join(outdir, "Native_HQ_GE"), ID=hq_id, model=model, chain=chain_id)
+        logger.info(f'High-quality native entanglements saved to {HQNativeEnt["outfile"]}')
+
+        # Cluster the native entanglements to remove degeneracies
+        nativeClusteredEnt = clustering.Cluster_NativeEntanglements(HQNativeEnt['outfile'], outdir=os.path.join(outdir, "Native_clustered_HQ_GE"), outfile=f"{hq_id}.csv", chain=chain_id)
+        logger.info(f'Clustered native entanglements saved to {nativeClusteredEnt["outfile"]}')
+
+        # Generate entanglement features for clustered native entanglements
+        FGen = FeatureGen(struct, outdir=os.path.join(outdir, "Native_clustered_HQ_GE_features"), cluster_file=nativeClusteredEnt['outfile'], log_level=log_level, logdir=logdir)
+        EntFeatures = FGen.get_uent_features(gene=args.Accession, chain=chain_id, pdbid=ID)
+        logger.info(f'Entanglement features saved to {EntFeatures["outfile"]}')
+
+
+    logger.info(f'NORMAL TERMINATION - {time.time() - start_time:.1f} seconds')
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())