# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# + [markdown] pycharm={"name": "#%% md\n"}
# # PyGromos File Examples

# + pycharm={"name": "#%%\n"}
import os, sys
root_dir = os.getcwd()

#if package is not installed and path not set correct - this helps you out :)
sys.path.append(root_dir+"/..")

# + [markdown] pycharm={"name": "#%% md\n"}
# ## IMD - Simulation Paramter File

# + pycharm={"name": "#%%\n"}
from pygromos import files

in_imd = "../pygromos/data/imd_templates/md.imd"
out_imd="test.imd"

#load
imd = files.Imd(in_imd)

#change Solvent number in system:
imd.SYSTEM.NSM = 500

#add EDS block
imd.edit_EDS(NUMSTATES=2, S=1.0, EIR=[0.0 for x in range(2)])

#get_file content:
print(imd)

#store again
out_imd = imd.write(out_imd)

# + [markdown] pycharm={"name": "#%% md\n"}
# ## CNF - Coordinate File

# + pycharm={"name": "#%%\n"}
"""
    CNF
"""
from pygromos import files

#load_cnf
in_cnf="in_dir/to/my.cnf"
out_cnf="in_dir/to/new.cnf"

cnf = files.Cnf(in_cnf)

#get residues of cnf
residues = cnf.get_residues()

#Delete Residues
delete_ligands = ["delete ME!"]
for resn in delete_ligands:
    cnf.delete_residue(resName=resn)

#set new title
cnf._blocks["TITLE"].content = "  Ligands:\t " + " ".join(lig_sys) + "\n"

#cleaning
cnf.clean_posiResNums()

#get_file content:
print(cnf)

#store again
out_cnf = cnf.write(out_path=out_cnf)
# -

# ## TRC - Coordinate Trajectory File

# + pycharm={"name": "#%%\n"}
#This snippet converts multiple cnfs to a trajectory.
import glob
from pygromos import files

in_cnfs_reg = "*.cnf"
out_trc_path = "concat_cnfs.trc"
cnfs = glob.glob(in_cnfs_reg)
out_trc = files.Trc.cnfs_to_trc(cnfs)
print(out_trc)
out_trc.write(out_trc_path)


# + [markdown] pycharm={"name": "#%% md\n"}
# ## TOP - topology File
#

# + [markdown] pycharm={"name": "#%% md\n"}
# ## Other Files
#
# ### MTB - topology building block file
#
# ### IFP - topology parameter file

# + pycharm={"name": "#%%\n"}
ifp_path = root+"\\examples\\test\\54a7.ifp"

sys.path.append(root+"/../../")

import pygromos
from pygromos.files.topology import ifp
#parse forcefield file
myfp = ifp.ifp(ifp_path)
myfp.write(os.getcwd()+"/fun.ifp")

# + pycharm={"name": "#%%\n"}
#parse output and write out again
test = ifp.ifp(os.getcwd()+"/fun.ifp")

test.write(os.getcwd()+"/fun.ifp")

# + [markdown] pycharm={"name": "#%% md\n"}
# ### disres - distance restraint file
# -

# ## PTP-Files in PyGromos
#
# Here I try to give a few example on what s possible with the ptp obj in pygromos.

# + pycharm={"name": "#%%\n"}
from pygromos.files.topology.ptp import Pertubation_topology as PTP
from pygromos.files.blocks.topology_blocks import  MPERTATOM, atom_eds_pertubation_state, pertubation_eds_state

# -

# ### defining some state types for later use :)

# + pycharm={"name": "#%%\n"}
dummy_type = pertubation_eds_state(IAC=22, CHARGE=0.0)
my_type = pertubation_eds_state(IAC=99, CHARGE=-1.0)
# -

# ### Read in ptp file:

# + pycharm={"name": "#%%\n"}
#Read in ptp file:
path= "../pygromos/tests/testfiles/ptp/eds_short.ptp"
ptp = PTP(path)
print(ptp)

# -

# ### Add atom or state or overwrite atominformation (except atom.NR)

# + pycharm={"name": "#%%\n"}
new_atoms_state = [atom_eds_pertubation_state(NR=x, NAME="H", STATES={7: my_type}) for x in range(1, 4)]
ptp.add_block(bloc=MPERTATOM(new_atoms_state))

print(ptp)
# -

# ### delete full state

# + pycharm={"name": "#%%\n"}
ptp.MPERTATOM.delete_state(stateIDs=[1,3])

print(ptp)
# -

# ### delete specific atoms

# + pycharm={"name": "#%%\n"}
ptp.MPERTATOM.delete_atom(atomNR=[1,2,7,8,9])
print(ptp)
# -

# #3# Write out ptp file

# + pycharm={"name": "#%%\n"}
ptp.write("fun.ptp")
# -

# ### Building ptp from scratch and generate all possible state combinations:

# + pycharm={"name": "#%%\n"}
print(ptp.MPERTATOM.STATEATOMHEADER)

# + pycharm={"name": "#%%\n"}
import numpy as np
from itertools import combinations
from pygromos.files.topology.ptp import Pertubation_topology as PTP
from pygromos.files.blocks.topology_blocks import  atom_eds_pertubation_state, pertubation_eds_state

#INPUT:
## parameters
outpath_new_ptp = "fun.ptp"

## the states to use:
o_state = pertubation_eds_state(IAC=16, CHARGE=-1.0)
h_state = pertubation_eds_state(IAC=33, CHARGE=-2.0)

## Map for molecule ID with Atom IDS
### First atom assumed to be O and last two Hs
molecules_atoms = {1: [1,2,3],
             2: [4,5,6],
             3: [7,8,9],}


#BUILD UP STATES
## Generate active state mapping:
max_active_mols_same_time = len(molecules_atoms)
molecule_states={}
state_ID=1
for active_mols in range(1, max_active_mols_same_time+1):
    combis = list(combinations(molecules_atoms, active_mols))
    molecule_states.update({state_ID+ind: x for ind, x in enumerate(combis)})
    state_ID = state_ID+len(combis)
#gives the state number as key and all the active molecules in this state
print("gives the state number as key and all the active molecules in this state")
print(molecule_states)
print()

#build state atoms for ptp
new_atoms_state_dict = {}
for state, molecules in molecule_states.items():
    for molecule in molecules:
        if(molecule in new_atoms_state_dict):
            atoms = new_atoms_state_dict[molecule]
            atoms[0].STATES.update({state: o_state})
            atoms[1].STATES.update({state: h_state})
            atoms[2].STATES.update({state: h_state})
        else:
            atoms = [atom_eds_pertubation_state(NR=molecules_atoms[molecule][0], NAME="O", STATES={state: o_state}),
                     atom_eds_pertubation_state(NR=molecules_atoms[molecule][1], NAME="H1", STATES={state: h_state}) ,
                     atom_eds_pertubation_state(NR=molecules_atoms[molecule][2], NAME="H2", STATES={state: h_state})]
            new_atoms_state_dict.update({molecule: atoms})

print("gives the atom_perttubation states for all mols")
print(new_atoms_state_dict)
print()

##finally make a list for our ptp file (#ThanksGromos)
new_atoms_state = np.concatenate(list(new_atoms_state_dict.values()))
#print(list(map(lambda x: x.STATES, new_atoms_state)))

#BUILD PTP
ptp = PTP()
ptp.MPERTATOM.add_state_atoms(state_atoms=new_atoms_state)
print(ptp)
ptp.write(outpath_new_ptp)

#TADAAAAA - DONE

StarcoderdataJupyterScriptsDedupFiltered