servalcat 0.4.60__cp38-cp38-win_amd64.whl → 0.4.72__cp38-cp38-win_amd64.whl

servalcat/__init__.py

@@ -6,5 +6,5 @@ This software is released under the
 Mozilla Public License, version 2.0; see LICENSE.
 """
 
-__version__ = '0.4.60'
-__date__ = '2024-02-26'
+__version__ = '0.4.72'
+__date__ = '2024-06-06'

servalcat/refine/refine.py

@@ -10,13 +10,13 @@ import os
 import re
 import gemmi
 import numpy
+import json
 import pandas
 import scipy.sparse
 import servalcat # for version
 from servalcat.utils import logger
 from servalcat import utils
 from servalcat.refmac import exte
-from servalcat.refmac.refmac_keywords import parse_keywords
 from servalcat import ext
 from . import cgsolve
 u_to_b = utils.model.u_to_b
@@ -29,7 +29,7 @@ b_to_u = utils.model.b_to_u
 
 class Geom:
     def __init__(self, st, topo, monlib, adpr_w=1, shake_rms=0,
-                 refmac_keywords=None, unrestrained=False, use_nucleus=False,
+                 params=None, unrestrained=False, use_nucleus=False,
                  ncslist=None, atom_pos=None):
         self.st = st
         self.atoms = [None for _ in range(self.st[0].count_atom_sites())]
@@ -48,26 +48,29 @@ class Geom:
             n_sym = len(images) + 1
             self.geom.specials.append(ext.Geometry.Special(atom, matp, mata, n_sym))
         self.adpr_w = adpr_w
+        self.occr_w = 1.
         self.unrestrained = unrestrained
         if shake_rms > 0:
             numpy.random.seed(0)
             utils.model.shake_structure(self.st, shake_rms, copy=False)
-            utils.fileio.write_model(self.st, "shaken", pdb=True, cif=True)
+            #utils.fileio.write_model(self.st, "shaken", pdb=True, cif=True)
+        self.use_nucleus = use_nucleus
+        self.calc_kwds = {"use_nucleus": self.use_nucleus}
+        if params is None:
+            params = {}
+        exte.read_external_restraints(params.get("exte", []), self.st, self.geom)
+        for k in ("wbond", "wangle", "wtors", "wplane", "wchir", "wvdw", "wncs"):
+            if k in params:
+                self.calc_kwds[k] = params[k]
+                logger.writeln("setting geometry weight {}= {}".format(k, params[k]))
+        inc_tors, exc_tors = utils.restraints.make_torsion_rules(params.get("restr", {}))
+        rtors = utils.restraints.select_restrained_torsions(monlib, inc_tors, exc_tors)
+        self.geom.mon_tors_names = rtors["monomer"]
+        self.geom.link_tors_names = rtors["link"]
+        self.group_occ = GroupOccupancy(self.st, params.get("occu"))
         if not self.unrestrained:
             self.geom.load_topo(topo)
             self.check_chemtypes(os.path.join(monlib.path(), "ener_lib.cif"), topo)
-        self.use_nucleus = use_nucleus
-        self.calc_kwds = {"use_nucleus": self.use_nucleus}
-        if refmac_keywords:
-            exte.read_external_restraints(refmac_keywords, self.st, self.geom)
-            kwds = parse_keywords(refmac_keywords)
-            for k in ("wbond", "wangle", "wtors", "wplane", "wchir", "wvdw", "wncs"):
-                if k in kwds:
-                    self.calc_kwds[k] = kwds[k]
-                    logger.writeln("setting geometry weight {}= {}".format(k, kwds[k]))
-            self.group_occ = GroupOccupancy(self.st, kwds.get("occu"))
-        else:
-            self.group_occ = GroupOccupancy(self.st, None)
         self.geom.finalize_restraints()
         self.outlier_sigmas = dict(bond=5, angle=5, torsion=5, vdw=5, ncs=5, chir=5, plane=5, staca=5, stacd=5, per_atom=5)
         self.parents = {}
@@ -104,20 +107,28 @@ class Geom:
         return self.geom.calc(check_only=target_only, **self.calc_kwds)
     def calc_adp_restraint(self, target_only):
         return self.geom.calc_adp_restraint(target_only, self.adpr_w)
-    def calc_target(self, target_only, refine_xyz, adp_mode):
+    def calc_occ_restraint(self, target_only):
+        return self.geom.calc_occ_restraint(target_only, self.occr_w)
+    def calc_target(self, target_only, refine_xyz, adp_mode, use_occr):
         self.geom.clear_target()
         geom_x = self.calc(target_only) if refine_xyz else 0
         geom_a = self.calc_adp_restraint(target_only) if adp_mode > 0 else 0
+        geom_q = self.calc_occ_restraint(target_only) if use_occr > 0 else 0
         logger.writeln(" geom_x = {}".format(geom_x))
         logger.writeln(" geom_a = {}".format(geom_a))
-        geom = geom_x + geom_a
+        logger.writeln(" geom_q = {}".format(geom_q))
+        geom = geom_x + geom_a + geom_q
         if not target_only:
             self.geom.spec_correction()
         return geom
         
-    def show_model_stats(self, show_outliers=True):
-        f0_x = self.calc(True)
-        f0_a = self.calc_adp_restraint(True)
+    def show_model_stats(self, refine_xyz=True, adp_mode=1, use_occr=False, show_outliers=True):
+        if refine_xyz:
+            self.calc(True)
+        if adp_mode > 0:
+            self.calc_adp_restraint(True)
+        if use_occr:
+            self.calc_occ_restraint(True)
         ret = {"outliers": {}}
         if show_outliers:
             get_table = dict(bond=self.geom.reporting.get_bond_outliers,
@@ -184,10 +195,12 @@ class Geom:
 def show_binstats(df, cycle_number):
     forplot = []
     rlabs = [x for x in df if x.startswith("R")]
+    fsclabs = [x for x in df if x.startswith("fsc")]
     cclabs = [x for x in df if x.startswith("CC")]
     dlabs = [x for x in df if re.search("^D[0-9]*", x)]
     if "fsc_model" in df: forplot.append(["FSC", ["fsc_model"]])
     if rlabs: forplot.append(["R", rlabs])
+    if fsclabs: forplot.append(["FSC", fsclabs])
     if cclabs: forplot.append(["CC", cclabs])
     if dlabs: forplot.append(["ML parameters - D", dlabs])
     if "S" in df: forplot.append(["ML parameters - Sigma", ["S"]])
@@ -197,6 +210,29 @@ def show_binstats(df, cycle_number):
     logger.writeln(lstr)
 # show_binstats()
 
+def convert_stats_to_dicts(stats):
+    tmp = []
+    for s in stats: # stats must be a list of dict
+        tmp.append({})
+        for k in s:
+            if k == "geom":
+                tmp[-1]["geom"] = {"summary": s["geom"]["summary"].to_dict()}
+                for kk in s["geom"]["outliers"]:
+                    tmp[-1]["geom"].setdefault("outliers", {})[kk] = s["geom"]["outliers"][kk].to_dict(orient="records")
+            else:
+                tmp[-1][k] = s[k]
+    return tmp
+# convert_stats_to_dicts()
+
+def write_stats_json_safe(stats, json_out):
+    tmp = convert_stats_to_dicts(stats)
+    out_tmp = json_out + ".part"
+    with open(out_tmp, "w") as ofs:
+        json.dump(tmp, ofs, indent=2)
+    os.replace(out_tmp, json_out)
+    logger.writeln(f"Refinement statistics saved: {json_out}")
+# write_stats_json_safe()
+
 class GroupOccupancy:
     # TODO max may not be one. should check multiplicity
     def __init__(self, st, params):
@@ -377,10 +413,11 @@ class GroupOccupancy:
 
 class Refine:
     def __init__(self, st, geom, ll=None, refine_xyz=True, adp_mode=1, refine_h=False, refine_occ=False,
-                 unrestrained=False, refmac_keywords=None):
+                 unrestrained=False, params=None):
         assert adp_mode in (0, 1, 2) # 0=fix, 1=iso, 2=aniso
         assert geom is not None
         self.st = st # clone()?
+        self.st_traj = None
         self.atoms = geom.atoms # not a copy
         self.geom = geom
         self.ll = ll
@@ -388,11 +425,15 @@ class Refine:
         self.adp_mode = 0 if self.ll is None else adp_mode
         self.refine_xyz = refine_xyz
         self.refine_occ = refine_occ
+        self.use_occr = self.refine_occ # for now?
         self.unrestrained = unrestrained
         self.refine_h = refine_h
         self.h_inherit_parent_adp = self.adp_mode > 0 and not self.refine_h and self.st[0].has_hydrogen()
         if self.h_inherit_parent_adp:
             self.geom.set_h_parents()
+        if params and params.get("write_trajectory"):
+            self.st_traj = self.st.clone()
+            self.st_traj[-1].name = "0"
         assert self.geom.group_occ.groups or self.n_params() > 0
     # __init__()
     
@@ -513,7 +554,7 @@ class Refine:
             self.ll.update_fc()
         
         self.geom.setup_nonbonded(self.refine_xyz) # if refine_xyz=False, no need to do it every time
-        self.geom.geom.setup_target(self.refine_xyz, self.adp_mode, self.refine_occ)
+        self.geom.geom.setup_target(self.refine_xyz, self.adp_mode, self.refine_occ, self.use_occr)
         logger.writeln("vdws = {}".format(len(self.geom.geom.vdws)))
 
     def get_x(self):
@@ -540,7 +581,7 @@ class Refine:
         N = self.n_params()
         geom = self.geom.calc_target(target_only,
                                      not self.unrestrained and self.refine_xyz,
-                                     self.adp_mode)
+                                     self.adp_mode, self.use_occr)
         if self.ll is not None:
             ll = self.ll.calc_target()
             logger.writeln(" ll= {}".format(ll))
@@ -627,14 +668,17 @@ class Refine:
 
         return ret, shift_scale, f1
 
-    def run_cycles(self, ncycles, weight=1, debug=False):
+    def run_cycles(self, ncycles, weight=1, weight_adjust=False, debug=False,
+                   weight_adjust_bond_rmsz_range=(0.5, 1.), stats_json_out=None):
         self.print_weights()
         stats = [{"Ncyc": 0}]
         self.geom.setup_nonbonded(self.refine_xyz)
-        self.geom.geom.setup_target(self.refine_xyz, self.adp_mode, self.refine_occ)
+        self.geom.geom.setup_target(self.refine_xyz, self.adp_mode, self.refine_occ, self.use_occr)
         logger.writeln("vdws = {}".format(len(self.geom.geom.vdws)))
-        if self.refine_xyz and not self.unrestrained:
-            stats[-1]["geom"] = self.geom.show_model_stats(show_outliers=True)["summary"]
+        stats[-1]["geom"] = self.geom.show_model_stats(refine_xyz=self.refine_xyz and not self.unrestrained,
+                                                       adp_mode=self.adp_mode,
+                                                       use_occr=self.refine_occ,
+                                                       show_outliers=True)
         if self.ll is not None:
             self.ll.update_fc()
             self.ll.overall_scale()
@@ -645,18 +689,26 @@ class Refine:
             show_binstats(llstats["bin_stats"], 0)
         if self.adp_mode > 0:
             utils.model.adp_analysis(self.st)
+        if stats_json_out:
+            write_stats_json_safe(stats, stats_json_out)
         occ_refine_flag = self.ll is not None and self.geom.group_occ.groups and self.geom.group_occ.ncycle > 0
 
         for i in range(ncycles):
             logger.writeln("\n====== CYCLE {:2d} ======\n".format(i+1))
-            if self.refine_xyz or self.adp_mode > 0:
+            logger.writeln(f" weight = {weight:.4e}")
+            if self.refine_xyz or self.adp_mode > 0 or self.refine_occ:
                 is_ok, shift_scale, fval = self.run_cycle(weight=weight)
-                stats.append({"Ncyc": len(stats), "shift_scale": shift_scale, "fval": fval, "fval_decreased": is_ok})
+                stats.append({"Ncyc": len(stats), "shift_scale": shift_scale, "fval": fval, "fval_decreased": is_ok,
+                              "weight": weight})
+            elif occ_refine_flag:
+                stats.append({"Ncyc": len(stats)})
             if occ_refine_flag:
                 stats[-1]["occ_refine"] = self.geom.group_occ.refine(self.ll, self.refine_h)
             if debug: utils.fileio.write_model(self.st, "refined_{:02d}".format(i+1), pdb=True)#, cif=True)
-            if self.refine_xyz and not self.unrestrained:
-                stats[-1]["geom"] = self.geom.show_model_stats(show_outliers=(i==ncycles-1))["summary"]
+            stats[-1]["geom"] = self.geom.show_model_stats(refine_xyz=self.refine_xyz and not self.unrestrained,
+                                                           adp_mode=self.adp_mode,
+                                                           use_occr=self.refine_occ,
+                                                           show_outliers=(i==ncycles-1))
             if self.ll is not None:
                 self.ll.overall_scale()
                 f0 = self.ll.calc_target()
@@ -670,6 +722,22 @@ class Refine:
                 show_binstats(llstats["bin_stats"], i+1)
             if self.adp_mode > 0:
                 utils.model.adp_analysis(self.st)
+            if (weight_adjust and self.refine_xyz and not self.unrestrained and self.ll is not None and
+                len(stats) > 2 and "Bond distances, non H" in stats[-1]["geom"]["summary"].index):
+                rmsz = stats[-1]["geom"]["summary"]["r.m.s.Z"]["Bond distances, non H"]
+                rmsz0 = stats[-2]["geom"]["summary"]["r.m.s.Z"]["Bond distances, non H"]
+                if rmsz > weight_adjust_bond_rmsz_range[1] and rmsz > rmsz0:
+                    weight /= 1.1
+                elif rmsz < weight_adjust_bond_rmsz_range[0] and rmsz0 < weight_adjust_bond_rmsz_range[0] and rmsz < rmsz0:
+                    weight *= 1.3
+                elif rmsz > 1.5 * rmsz0:
+                    weight /= 1.1
+            if self.st_traj is not None:
+                self.st_traj.add_model(self.st[0])
+                self.st_traj[-1].name = str(i+1)
+            if stats_json_out:
+                write_stats_json_safe(stats, stats_json_out)
+
             logger.writeln("")
 
         # Make table
@@ -685,8 +753,8 @@ class Refine:
                                 ("r.m.s.Z", "Bond distances, non H", "zBOND"),
                                 ("r.m.s.d.", "Bond angles, non H", "rmsANGL"),
                                 ("r.m.s.Z", "Bond angles, non H", "zANGL")):
-                    if k in d["geom"] and n in d["geom"][k]:
-                        x[l] = d["geom"][k].get(n)
+                    if k in d["geom"]["summary"] and n in d["geom"]["summary"][k]:
+                        x[l] = d["geom"]["summary"][k].get(n)
                         geom_keys.add(l)
             tmp.append(x)
         df = pandas.DataFrame(tmp)
@@ -711,12 +779,20 @@ class Refine:
         lstr = utils.make_loggraph_str(df, "stats vs cycle", forplot,
                                        float_format="{:.4f}".format)
         logger.writeln(lstr)
-        self.update_meta()
+        self.update_meta(stats[-1])
         return stats
 
-    def update_meta(self):
+    def update_meta(self, stats):
         # TODO write stats. probably geom.reporting.get_summary_table should return with _refine_ls_restr.type names
-        self.st.raw_remarks = []
+        # should remove st.mod_residues?
+        self.st.helices.clear()
+        self.st.sheets.clear()
+        raw_remarks = [f'REMARK   3',
+                       f'REMARK   3 REFINEMENT.',
+                       f'REMARK   3   PROGRAM     : SERVALCAT {servalcat.__version__}',
+                       f'REMARK   3   AUTHORS     : YAMASHITA,MURSHUDOV',
+                       f'REMARK   3',
+                       ]
         si = gemmi.SoftwareItem()
         si.classification = gemmi.SoftwareItem.Classification.Refinement
         si.name = "Servalcat"
@@ -724,10 +800,22 @@ class Refine:
         si.date = servalcat.__date__
         self.st.meta.software = [si]
 
-        self.st.meta.refinement = []
-        #ri = gemmi.RefinementInfo()
-        #rr = gemmi.RefinementInfo.Restr("")
-        #ri.restr_stats.append(rr)
-        #st.meta.refinement = [ri]
-        
+        ri = gemmi.RefinementInfo()
+        if "geom" in stats:
+            restr_stats = []
+            raw_remarks.append("REMARK   3  RMS DEVIATIONS FROM IDEAL VALUES        COUNT    RMS    WEIGHT")
+            for k, n, l, pl in (("r.m.s.d.", "Bond distances, non H", "s_bond_nonh_d", "BOND LENGTHS REFINED ATOMS        (A)"),
+                                ("r.m.s.d.", "Bond angles, non H", "s_angle_nonh_d", "BOND ANGLES REFINED ATOMS   (DEGREES)")):
+                if k in stats["geom"]["summary"] and n in stats["geom"]["summary"][k]:
+                    rr = gemmi.RefinementInfo.Restr(l)
+                    rr.dev_ideal = stats["geom"]["summary"][k].get(n)
+                    rr.count = stats["geom"]["summary"]["N restraints"].get(n)
+                    rr.weight = stats["geom"]["summary"]["Mn(sigma)"].get(n)
+                    restr_stats.append(rr)
+                    raw_remarks.append(f"REMARK   3   {pl}:{rr.count:6d} ;{rr.dev_ideal:6.3f} ;{rr.weight:6.3f}")
+            ri.restr_stats = restr_stats
+            raw_remarks.append("REMARK   3")
+        self.st.meta.refinement = [ri]
+        self.st.raw_remarks = raw_remarks
+
 # class Refine

servalcat/refine/refine_geom.py

@@ -14,7 +14,8 @@ import json
 import servalcat # for version
 from servalcat.utils import logger
 from servalcat import utils
-from servalcat.refine.refine import Geom, Refine
+from servalcat.refine.refine import Geom, Refine, convert_stats_to_dicts
+from servalcat.refmac import refmac_keywords
 
 def add_arguments(parser):
     group = parser.add_mutually_exclusive_group(required=True)
@@ -53,24 +54,32 @@ def parse_args(arg_list):
 # parse_args()
 
 def add_program_info_to_dictionary(block, comp_id, program_name="servalcat", descriptor="optimization tool"):
-    tab = block.find("_pdbx_chem_comp_description_generator.", ["program_name", "program_version", "descriptor"])
-    # just overwrite version if it's there
-    for row in tab:
-        if row.str(0) == program_name and row.str(2) == descriptor:
-            row[1] = gemmi.cif.quote(servalcat.__version__)
-            return
-    loop = tab.loop
-    if not loop:
-        loop = block.init_loop("_pdbx_chem_comp_description_generator.", ["comp_id",
-                                                                          "program_name",
-                                                                          "program_version",
-                                                                          "descriptor"])
+    # old acedrg used _pdbx_chem_comp_description_generator. and descriptor
+    # new acedrg (>280?) uses _acedrg_chem_comp_descriptor. and type
+    for tag, name in (("_acedrg_chem_comp_descriptor.", "type"),
+                      ("_pdbx_chem_comp_description_generator.", "descriptor")):
+        tab = block.find(tag, ["program_name", "program_version", name])
+        if tab:
+            loop = tab.loop
+            # just overwrite version if it's there
+            for row in tab:
+                if row.str(0) == program_name and row.str(2) == descriptor:
+                    row[1] = gemmi.cif.quote(servalcat.__version__)
+                    return
+            break
+    else:
+        # it may be strange to say _acedrg in this case..
+        name = "type"
+        loop = block.init_loop("_acedrg_chem_comp_descriptor.", ["comp_id",
+                                                                 "program_name",
+                                                                 "program_version",
+                                                                 name])
     tags = [x[x.index(".")+1:] for x in loop.tags]
     row = ["" for _ in range(len(tags))]
     for tag, val in (("comp_id", comp_id),
                      ("program_name", program_name),
                      ("program_version", servalcat.__version__),
-                     ("descriptor", descriptor)):
+                     (name, descriptor)):
         if tag in tags: row[tags.index(tag)] = val
     loop.add_row(gemmi.cif.quote_list(row))
 # add_program_info_to_dictionary()
@@ -122,18 +131,16 @@ def refine_and_update_dictionary(cif_in, monomer_dir, output_prefix, randomize=0
     doc.write_file(output_prefix + "_updated.cif", style=gemmi.cif.Style.Aligned)
     logger.writeln("Updated dictionary saved: {}".format(output_prefix + "_updated.cif"))
     with open(output_prefix + "_stats.json", "w") as ofs:
-        for stats in all_stats:
-            for s in stats:
-                s["geom"] = s["geom"].to_dict()
-        json.dump(all_stats, ofs, indent=2)
+        json.dump([convert_stats_to_dicts(x) for x in all_stats],
+                  ofs, indent=2)
         logger.writeln("Refinement statistics saved: {}".format(ofs.name))
 # refine_and_update_dictionary()
 
-def refine_geom(model_in, monomer_dir, cif_files, h_change, ncycle, output_prefix, randomize, refmac_keywords,
+def refine_geom(model_in, monomer_dir, cif_files, h_change, ncycle, output_prefix, randomize, params,
                 find_links=False, use_ncsr=False):
     st = utils.fileio.read_structure(model_in)
     utils.model.setup_entities(st, clear=True, force_subchain_names=True, overwrite_entity_type=True)
-    if st.ncs:
+    if not all(op.given for op in st.ncs):
         st2 = st.clone()
         logger.writeln("Take NCS constraints into account.")
         st2.expand_ncs(gemmi.HowToNameCopiedChain.Dup, merge_dist=0)
@@ -141,27 +148,26 @@ def refine_geom(model_in, monomer_dir, cif_files, h_change, ncycle, output_prefi
 
     monlib = utils.restraints.load_monomer_library(st, monomer_dir=monomer_dir,
                                                    cif_files=cif_files,
-                                                   stop_for_unknowns=True)
+                                                   stop_for_unknowns=True,
+                                                   params=params)
     utils.restraints.find_and_fix_links(st, monlib, add_found=find_links) # should remove unknown id here?
     try:
-        topo, metal_kws = utils.restraints.prepare_topology(st, monlib, h_change=h_change,
-                                                            check_hydrogen=(h_change==gemmi.HydrogenChange.NoChange))
+        topo, _ = utils.restraints.prepare_topology(st, monlib, h_change=h_change,
+                                                    check_hydrogen=(h_change==gemmi.HydrogenChange.NoChange),
+                                                    params=params)
     except RuntimeError as e:
         raise SystemExit("Error: {}".format(e))
-    refmac_keywords = metal_kws + refmac_keywords
+
     if use_ncsr:
         ncslist = utils.restraints.prepare_ncs_restraints(st)
     else:
         ncslist = False
-    geom = Geom(st, topo, monlib, shake_rms=randomize, refmac_keywords=refmac_keywords, ncslist=ncslist)
+    geom = Geom(st, topo, monlib, shake_rms=randomize, params=params, ncslist=ncslist)
     refiner = Refine(st, geom)
-    stats = refiner.run_cycles(ncycle)
+    stats = refiner.run_cycles(ncycle,
+                               stats_json_out=output_prefix + "_stats.json")
     refiner.st.name = output_prefix
     utils.fileio.write_model(refiner.st, output_prefix, pdb=True, cif=True)
-    with open(output_prefix + "_stats.json", "w") as ofs:
-        for s in stats: s["geom"] = s["geom"].to_dict()
-        json.dump(stats, ofs, indent=2)
-        logger.writeln("Refinement statistics saved: {}".format(ofs.name))
 # refine_geom()
 
 def main(args):
@@ -170,6 +176,7 @@ def main(args):
     if args.keyword_file: keywords.extend(l for f in sum(args.keyword_file, []) for l in open(f))
     decide_prefix = lambda f: utils.fileio.splitext(os.path.basename(f))[0] + "_refined"
     if args.model:
+        params = refmac_keywords.parse_keywords(keywords)
         if not args.output_prefix:
             args.output_prefix = decide_prefix(args.model)
         if args.ligand:
@@ -185,7 +192,7 @@ def main(args):
                     ncycle=args.ncycle,
                     output_prefix=args.output_prefix,
                     randomize=args.randomize,
-                    refmac_keywords=keywords,
+                    params=params,
                     find_links=args.find_links,
                     use_ncsr=args.ncsr)
     else:

servalcat/refine/refine_spa.py

@@ -8,7 +8,6 @@ Mozilla Public License, version 2.0; see LICENSE.
 from __future__ import absolute_import, division, print_function, generators
 import gemmi
 import argparse
-import json
 import numpy
 from servalcat.utils import logger
 from servalcat import utils
@@ -16,6 +15,7 @@ from servalcat.spa.run_refmac import check_args, process_input, calc_fsc, calc_f
 from servalcat.spa import fofc
 from servalcat.refine import spa
 from servalcat.refine.refine import Geom, Refine
+from servalcat.refmac import refmac_keywords
 b_to_u = utils.model.b_to_u
 
 def add_arguments(parser):
@@ -56,6 +56,7 @@ def add_arguments(parser):
     parser.add_argument('--hydrogen', default="all", choices=["all", "yes", "no"],
                         help="all: add riding hydrogen atoms, yes: use hydrogen atoms if present, no: remove hydrogen atoms in input. "
                         "Default: %(default)s")
+    parser.add_argument('--hout', action='store_true', help="write hydrogen atoms in the output model")
     parser.add_argument('--jellybody', action='store_true',
                         help="Use jelly body restraints")
     parser.add_argument('--jellybody_params', nargs=2, type=float,
@@ -85,6 +86,10 @@ def add_arguments(parser):
                         help="number of CG cycles (default: %(default)d)")
     parser.add_argument('--weight', type=float,
                         help="refinement weight. default: automatic")
+    parser.add_argument('--no_weight_adjust', action='store_true', 
+                        help='Do not adjust weight during refinement')
+    parser.add_argument('--target_bond_rmsz_range', nargs=2, type=float, default=[0.5, 1.],
+                        help='Bond rmsz range for weight adjustment (default: %(default)s)')
     parser.add_argument('--adpr_weight', type=float, default=1.,
                         help="ADP restraint weight in B (default: %(default)f)")
     parser.add_argument('--ncsr', action='store_true', 
@@ -93,6 +98,7 @@ def add_arguments(parser):
                         help="reset all atomic B values to specified value")
     parser.add_argument('--fix_xyz', action="store_true")
     parser.add_argument('--adp',  choices=["fix", "iso", "aniso"], default="iso")
+    parser.add_argument('--refine_all_occ', action="store_true")
     parser.add_argument('--max_dist_for_adp_restraint', type=float, default=4.)
     parser.add_argument('--adp_restraint_power', type=float)
     parser.add_argument('--adp_restraint_exp_fac', type=float)
@@ -112,6 +118,8 @@ def add_arguments(parser):
                         help="Use scattering factor for charged atoms. Use it with care.")
     parser.add_argument("--keep_entities", action='store_true',
                         help="Do not override entities")
+    parser.add_argument("--write_trajectory", action='store_true',
+                        help="Write all output from cycles")
 # add_arguments()
 
 def parse_args(arg_list):
@@ -125,13 +133,14 @@ def main(args):
     args.invert_mask = False
     args.trim_fofc_mtz = args.mask_for_fofc is not None
     args.cross_validation_method = "throughout"
-    check_args(args)    
-    refmac_keywords = args.keywords + [l for f in args.keyword_file for l in open(f)]
+    check_args(args)
+    params = refmac_keywords.parse_keywords(args.keywords + [l for f in args.keyword_file for l in open(f)])
+    params["write_trajectory"] = args.write_trajectory
 
     st = utils.fileio.read_structure(args.model)
     try:
         monlib = utils.restraints.load_monomer_library(st, monomer_dir=args.monlib, cif_files=args.ligand,
-                                                       stop_for_unknowns=True)
+                                                       stop_for_unknowns=True, params=params)
     except RuntimeError as e:
         raise SystemExit("Error: {}".format(e))
     if not args.keep_entities:
@@ -165,14 +174,14 @@ def main(args):
                 "yes":gemmi.HydrogenChange.NoChange,
                 "no":gemmi.HydrogenChange.Remove}[args.hydrogen]
     try:
-        topo, metal_kws = utils.restraints.prepare_topology(st, monlib, h_change=h_change,
-                                                            check_hydrogen=(args.hydrogen=="yes"))
+        topo, _ = utils.restraints.prepare_topology(st, monlib, h_change=h_change,
+                                                    check_hydrogen=(args.hydrogen=="yes"),
+                                                    params=params)
     except RuntimeError as e:
         raise SystemExit("Error: {}".format(e))
-    refmac_keywords = metal_kws + refmac_keywords
+
     # initialize ADP
-    if args.adp != "fix":
-        utils.model.reset_adp(st[0], args.bfactor, args.adp == "aniso")
+    utils.model.reset_adp(st[0], args.bfactor, args.adp)
 
     # auto weight
     if args.weight is None:
@@ -200,7 +209,7 @@ def main(args):
     else:
         ncslist = False
     geom = Geom(st, topo, monlib, shake_rms=args.randomize, adpr_w=args.adpr_weight,
-                refmac_keywords=refmac_keywords, unrestrained=args.jellyonly,
+                params=params, unrestrained=args.jellyonly,
                 ncslist=ncslist)
     ll = spa.LL_SPA(hkldata, st, monlib,
                     lab_obs="F_map1" if args.cross_validation else "FP",
@@ -209,7 +218,8 @@ def main(args):
                      refine_xyz=not args.fix_xyz,
                      adp_mode=dict(fix=0, iso=1, aniso=2)[args.adp],
                      refine_h=args.refine_h,
-                     refmac_keywords=refmac_keywords)
+                     params=params,
+                     refine_occ=args.refine_all_occ)
 
     geom.geom.adpr_max_dist = args.max_dist_for_adp_restraint
     if args.adp_restraint_power is not None: geom.geom.adpr_d_power = args.adp_restraint_power
@@ -223,19 +233,18 @@ def main(args):
     #for cra in st[0].all():
     #    cra.atom.pos += gemmi.Position(0.3,0,0)
 
-    stats = refiner.run_cycles(args.ncycle, weight=args.weight)
+    stats = refiner.run_cycles(args.ncycle, weight=args.weight,
+                               weight_adjust=not args.no_weight_adjust,
+                               weight_adjust_bond_rmsz_range=args.target_bond_rmsz_range,
+                               stats_json_out=args.output_prefix + "_stats.json")
     if not args.hklin and not args.no_trim:
         refiner.st.cell = maps[0][0].unit_cell
         refiner.st.setup_cell_images()
 
     refiner.st.name = args.output_prefix
-    utils.fileio.write_model(refiner.st, args.output_prefix, pdb=True, cif=True)
-    with open(args.output_prefix + "_stats.json", "w") as ofs:
-        for s in stats:
-            if "geom" in s: s["geom"] = s["geom"].to_dict()
-        json.dump(stats, ofs, indent=2)
-        logger.writeln("Refinement statistics saved: {}".format(ofs.name))
-
+    utils.fileio.write_model(refiner.st, args.output_prefix, pdb=True, cif=True, hout=args.hout)
+    if params["write_trajectory"]:
+        utils.fileio.write_model(refiner.st_traj, args.output_prefix + "_traj", cif=True)
     if args.hklin:
         return
         
@@ -243,7 +252,7 @@ def main(args):
     st_expanded = refiner.st.clone()
     if not all(op.given for op in st.ncs):
         utils.model.expand_ncs(st_expanded)
-        utils.fileio.write_model(st_expanded, args.output_prefix+"_expanded", pdb=True, cif=True)
+        utils.fileio.write_model(st_expanded, args.output_prefix+"_expanded", pdb=True, cif=True, hout=args.hout)
 
     # Calc FSC
     mask = utils.fileio.read_ccp4_map(args.mask)[0] if args.mask else None
@@ -252,7 +261,7 @@ def main(args):
                            soft_edge=args.mask_soft_edge,
                            b_before_mask=args.b_before_mask,
                            no_sharpen_before_mask=args.no_sharpen_before_mask,
-                           make_hydrogen=args.hydrogen,
+                           make_hydrogen="yes", # no change needed in the model
                            monlib=monlib, 
                            blur=args.blur,
                            d_min_fsc=args.fsc_resolution,
@@ -273,9 +282,9 @@ def main(args):
         adpstats_txt += " Chain {0:{1}s}".format(chain, max_chain_len) if chain!="*" else " {0:{1}s}".format("All", max_chain_len+6)
         adpstats_txt += " ({0:{1}d} atoms) min={2:5.1f} median={3:5.1f} max={4:5.1f} A^2\n".format(natoms, max_num_len, qs[0],qs[2],qs[4])
 
-    if "geom" in stats[-1]:
-        rmsbond = stats[-1]["geom"]["r.m.s.d."]["Bond distances, non H"]
-        rmsangle = stats[-1]["geom"]["r.m.s.d."]["Bond angles, non H"]
+    if "geom" in stats[-1] and "Bond distances, non H" in stats[-1]["geom"]["summary"].index:
+        rmsbond = stats[-1]["geom"]["summary"]["r.m.s.d."]["Bond distances, non H"]
+        rmsangle = stats[-1]["geom"]["summary"]["r.m.s.d."]["Bond angles, non H"]
     else:
         rmsbond, rmsangle = numpy.nan, numpy.nan
     if args.mask_for_fofc: