servalcat 0.4.60__cp312-cp312-macosx_11_0_arm64.whl

servalcat/spa/fofc.py

@@ -0,0 +1,462 @@
+# TODO shift map first using mask!
+"""
+Author: "Keitaro Yamashita, Garib N. Murshudov"
+MRC Laboratory of Molecular Biology
+    
+This software is released under the
+Mozilla Public License, version 2.0; see LICENSE.
+"""
+from __future__ import absolute_import, division, print_function, generators
+import gemmi
+import numpy
+import time
+from servalcat.utils import logger
+from servalcat import utils
+from servalcat.spa import shift_maps
+import argparse
+
+def add_arguments(parser):
+    parser.description = 'Fo-Fc map calculation based on model and data errors'
+    group = parser.add_mutually_exclusive_group()
+    group.add_argument("--halfmaps", nargs=2)
+    group.add_argument("--map", help="Use only if you really do not have half maps.")
+    parser.add_argument('--pixel_size', type=float,
+                        help='Override pixel size (A)')
+    parser.add_argument('--model', required=True,
+                        help='Input atomic model file')
+    parser.add_argument("-d", '--resolution', type=float, required=True)
+    parser.add_argument('-m', '--mask', help="mask file")
+    parser.add_argument('-r', '--mask_radius', type=float, help="mask radius (not used if --mask is given)")
+    parser.add_argument('--no_check_mask_with_model', action='store_true', 
+                        help='Disable mask test using model')
+    parser.add_argument("-B", type=float, help="Estimated blurring")
+    parser.add_argument("--half1_only", action='store_true', help="Only use half 1 for map calculation (use half 2 only for noise estimation)")
+    parser.add_argument("--normalized_map", action='store_true',
+                        help="Write normalized map in the masked region. Now this is on by default.")
+    parser.add_argument("--no_fsc_weights", action='store_true',
+                        help="Just for debugging purpose: turn off FSC-based weighting")
+    parser.add_argument("--sharpening_b", type=float,
+                        help="Use B value (negative value for sharpening) instead of standard deviation of the signal")
+    parser.add_argument("--trim", action='store_true',
+                        help="Write trimmed maps")
+    parser.add_argument("--trim_mtz", action='store_true',
+                        help="Write trimmed mtz")
+    parser.add_argument("--monlib",
+                        help="Monomer library path. Default: $CLIBD_MON")
+    parser.add_argument("--omit_proton", action='store_true',
+                        help="Omit proton from model in map calculation")
+    parser.add_argument("--omit_h_electron", action='store_true',
+                        help="Omit hydrogen electrons from model in map calculation")
+    parser.add_argument("--source", choices=["electron", "xray", "neutron"], default="electron")
+    parser.add_argument('-o','--output_prefix', default="diffmap",
+                        help='output file name prefix (default: %(default)s)')
+    parser.add_argument('--keep_charges',  action='store_true',
+                        help="Use scattering factor for charged atoms. Use it with care.")
+# add_arguments()
+
+def parse_args(arg_list):
+    parser = argparse.ArgumentParser()
+    add_arguments(parser)
+    return parser.parse_args(arg_list)
+# parse_args()
+
+def calc_D_and_S(hkldata, has_halfmaps=True, half1_only=False):#fo_asu, fc_asu, varn, bins, bin_idxes):
+    bdf = hkldata.binned_df
+    bdf["D"] = 0.
+    bdf["S"] = 0.
+    stats_str = """$TABLE: Statistics :
+$GRAPHS
+: log(Mn(|F|^2)) and variances :A:1,6,7,8,13,14:
+: FSC :A:1,9,10,11:
+: weights :A:1,12,15,16:
+: map weights :A:1,17:
+$$
+1/resol^2 bin n d_max d_min log(var(Fo)) log(var(Fc)) log(var(DFc)) FSC.model FSC.full sqrt(FSC.full) D log(var_U,T) log(var_noise) wFo wFc wFo.sharpen
+$$
+$$
+"""
+    tmpl = "{:.4f} {:3d} {:7d} {:7.3f} {:7.3f} {:.4e} {:.4e} {:4e} {: .4f}   {: .4f} {: .4f} {: .4e} {:.4e} {:.4e} {:.4f} {:.4f} {:.4e}\n"
+
+    var_noise = None
+    FP = hkldata.df.FP.to_numpy()
+    if half1_only:
+        FP = hkldata.df.F_map1.to_numpy()
+        var_noise = hkldata.binned_df.var_noise * 2
+    elif has_halfmaps:
+        var_noise = hkldata.binned_df.var_noise
+        
+    for i_bin, idxes in hkldata.binned():
+        bin_d_min = hkldata.binned_df.d_min[i_bin]
+        bin_d_max = hkldata.binned_df.d_max[i_bin]
+        Fo = FP[idxes]
+        Fc = hkldata.df.FC.to_numpy()[idxes]
+        fsc = numpy.real(numpy.corrcoef(Fo, Fc)[1,0])
+        bdf.loc[i_bin, "D"] = numpy.sum(numpy.real(Fo * numpy.conj(Fc)))/numpy.sum(numpy.abs(Fc)**2)
+        if has_halfmaps:
+            varn = var_noise[i_bin]
+            fsc_full = hkldata.binned_df.FSCfull[i_bin]
+            S = max(0, numpy.average(numpy.abs(Fo-bdf.D[i_bin]*Fc)**2)-varn)
+            bdf.loc[i_bin, "S"] = S
+            w = S/(S+varn)
+            if fsc_full < 0: # this should be fixed actually. needs smoothing to zero.
+                w_sharpen = 0
+            else:
+                w_sharpen = w / numpy.sqrt(fsc_full) / numpy.std(Fo)
+        else:
+            varn = fsc_full = 0
+            w = 1
+            w_sharpen = 1
+
+        with numpy.errstate(divide="ignore", invalid="ignore"):
+            stats_str += tmpl.format(1/bin_d_min**2, i_bin, Fo.size, bin_d_max, bin_d_min,
+                                     numpy.log(numpy.average(numpy.abs(Fo)**2)),
+                                     numpy.log(numpy.average(numpy.abs(Fc)**2)),
+                                     numpy.log(bdf.D[i_bin]**2*numpy.average(numpy.abs(Fc)**2)),
+                                     fsc, fsc_full, numpy.sqrt(fsc_full), bdf.D[i_bin],
+                                     numpy.log(bdf.S[i_bin]), numpy.log(varn),
+                                     w, 1-w, w_sharpen)
+    return stats_str
+# calc_D_and_S()
+
+#import line_profiler
+#profile = line_profiler.LineProfiler()
+#import atexit
+#atexit.register(profile.print_stats)
+#@profile
+def calc_maps(hkldata, B=None, has_halfmaps=True, half1_only=False, no_fsc_weights=False, sharpening_b=None):
+    has_fc = "FC" in hkldata.df
+
+    if not has_fc:
+        labs = ["FWT"]
+        if B is not None: labs.append("FWT_b0")
+    elif has_halfmaps:
+        labs = ["Fupdate", "DELFWT", "FWT", "DELFWT_noscale", "Fupdate_noscale"]
+        if B is not None: labs.extend(["Fupdate_b0", "DELFWT_b0", "FWT_b0"])
+    else:
+        labs = ["DELFWT"]
+        
+    tmp = {}
+    for l in labs:
+        tmp[l] = numpy.zeros(len(hkldata.df.index), numpy.complex128)
+
+    logger.writeln("Calculating maps..")
+    logger.write(" sharpening method: ")
+    if sharpening_b is None:
+        logger.writeln("1/sqrt(FSC * Mn(Fo)) for Fo and 1/sigma_U,T for Fo-Fc")
+    else:
+        logger.writeln("1/exp(-B*s^2/4) with B= {:.2f}".format(sharpening_b))
+
+    time_t = time.time()
+
+    if half1_only:
+        FP = hkldata.df.F_map1.to_numpy()
+    else:
+        FP = hkldata.df.FP.to_numpy()
+
+    s2 = 1./hkldata.d_spacings().to_numpy()**2
+
+    fsc_became_negative = False
+        
+    for i_bin, idxes in hkldata.binned():
+        if has_halfmaps:
+            fsc = hkldata.binned_df.FSCfull[i_bin] # FSCfull
+            if half1_only:
+                varn = hkldata.binned_df.var_noise[i_bin] * 2
+                fsc = fsc/(2-fsc) # to FSChalf
+            else:
+                varn = hkldata.binned_df.var_noise[i_bin]
+        else:
+            fsc, varn = 1., 0.
+            
+        w_nomodel = 1. if no_fsc_weights else fsc
+        Fo = FP[idxes]
+        sig_fo = numpy.std(Fo)
+        s2_bin = s2[idxes]
+
+        if has_fc:
+            Fc = hkldata.df.FC.to_numpy()[idxes]
+            D = hkldata.binned_df.D[i_bin]
+            S = hkldata.binned_df.S[i_bin] # variance of unexplained signal
+            w = 1. if no_fsc_weights or not has_halfmaps else S/(S+varn)
+            delfwt = w * (Fo-D*Fc)
+            fup = 2 * w * Fo + (1 - 2*w) * D*Fc # <F> + delfwt
+            if has_halfmaps: # no point making this map when half maps not given
+                tmp["DELFWT_noscale"][idxes] = delfwt
+                tmp["Fupdate_noscale"][idxes] = fup
+
+        if not fsc_became_negative and fsc <= 0:
+            logger.writeln(" WARNING: cutting resolution at {:.2f} A because fsc < 0".format(hkldata.binned_df.d_max[i_bin]))
+            fsc_became_negative = True
+        if fsc_became_negative:
+            continue
+
+        if sharpening_b is None:
+            k = sig_fo * numpy.sqrt(fsc)
+            k_fofc = numpy.sqrt(S) if has_fc and S > 0 else 1. # to avoid zero-division. if S=0 then w=0.
+        else:
+            k = numpy.exp(-sharpening_b*s2_bin/4)
+            
+        lab_suf = "" if B is None else "_b0"
+        if has_halfmaps:
+            tmp["FWT"+lab_suf][idxes] = w_nomodel / k * Fo
+            if has_fc:
+                tmp["DELFWT"+lab_suf][idxes] = delfwt / k_fofc
+                tmp["Fupdate"+lab_suf][idxes] = fup / k
+        elif has_fc:
+            tmp["DELFWT"+lab_suf][idxes] = delfwt
+
+        if B is not None and has_halfmaps: # local B based map
+            k_l = numpy.exp(-B*s2_bin/4.)
+            k2_l = numpy.exp(-B*s2_bin/2.)
+            fsc_l = k2_l*fsc/(1+(k2_l-1)*fsc)
+            w_nomodel = 1. if no_fsc_weights else fsc_l
+            tmp["FWT"][idxes] = Fo*w_nomodel/k/k_l
+            if has_fc:
+                S_l = S * k2_l
+                w = 1. if no_fsc_weights or not has_halfmaps else S_l/(S_l+varn)            
+                delfwt = (Fo-D*Fc)*w/k_fofc/k_l
+                fup = (w*Fo+(1.-w)*D*Fc)/k/k_l
+                logger.writeln("{:4d} {:.4e} {:.4e} {:.4e} {:.4e} {:.4e} {:.4e}".format(i_bin,
+                                                                                      numpy.average(k),
+                                                                                      numpy.average(sig_fo),
+                                                                                      numpy.average(fsc_l),
+                                                                                      numpy.average(k_l),
+                                                                                      numpy.average(abs(fup)),
+                                                                                      numpy.average(abs(delfwt))))
+                tmp["DELFWT"][idxes] = delfwt
+                tmp["Fupdate"][idxes] = fup
+
+    for l in labs:
+        hkldata.df[l] = tmp[l]
+
+    logger.writeln(" finished in {:.3f} sec.".format(time.time()-time_t))
+    return labs
+# calc_maps()
+
+def dump_to_mtz(hkldata, map_labs, mtz_out):
+    extra_labs = list(filter(lambda x: x in hkldata.df, ["FP", "FC"]))
+    map_labs = map_labs + extra_labs
+    hkldata.write_mtz(mtz_out, map_labs)
+# dump_to_mtz()
+
+def calc_fofc(st, d_min, maps, mask=None, monlib=None, B=None, half1_only=False,
+              no_fsc_weights=False, sharpening_b=None, omit_proton=False, omit_h_electron=False,
+              source="electron"):
+    if no_fsc_weights:
+        logger.writeln("WARNING: --no_fsc_weights is requested.")
+    if sharpening_b is not None:
+        logger.writeln("WARNING: --sharpening_b={} is given".format(sharpening_b))
+    
+    hkldata = utils.maps.mask_and_fft_maps(maps, d_min, mask)
+    hkldata.df["FC"] = utils.model.calc_fc_fft(st, d_min - 1e-6, monlib=monlib, source=source,
+                                               miller_array=hkldata.miller_array())
+    if mask is not None:
+        fc_map = hkldata.fft_map("FC", grid_size=mask.shape)
+        fc_map.array[:] *= mask
+        hkldata.df["FC"] = gemmi.transform_map_to_f_phi(fc_map).get_value_by_hkl(hkldata.miller_array())
+        
+    hkldata.setup_relion_binning()
+
+    has_halfmaps = (len(maps) == 2)
+    if has_halfmaps:
+        utils.maps.calc_noise_var_from_halfmaps(hkldata)
+
+    stats_str = calc_D_and_S(hkldata, has_halfmaps=has_halfmaps, half1_only=half1_only)
+
+    if omit_proton or omit_h_electron:
+        hkldata.df["FC"] = utils.model.calc_fc_fft(st, d_min - 1e-6, monlib=monlib, source=source,
+                                                   omit_proton=omit_proton, omit_h_electron=omit_h_electron,
+                                                   miller_array=hkldata.miller_array())
+    
+    map_labs = calc_maps(hkldata, B=B, has_halfmaps=has_halfmaps, half1_only=half1_only,
+                         no_fsc_weights=no_fsc_weights, sharpening_b=sharpening_b)
+    return hkldata, map_labs, stats_str
+# calc_fofc()
+
+def write_files(hkldata, map_labs, grid_start, stats_str,
+                mask=None, output_prefix="diffmap", trim_map=False, trim_mtz=False,
+                normalize_map=True, omit_h_electron=False):
+    # this function may modify the overall scale of FWT/DELFWT.
+
+    if mask is not None and (trim_map or trim_mtz):
+        new_cell, new_shape, new_grid_start, shifts = shift_maps.determine_shape_and_shift(mask=gemmi.FloatGrid(mask.array,
+                                                                                                                hkldata.cell,
+                                                                                                                hkldata.sg),
+                                                                                           grid_start=grid_start,
+                                                                                           padding=5,
+                                                                                           mask_cutoff=0.5,
+                                                                                           noncentered=True,
+                                                                                           noncubic=True,
+                                                                                           json_out=None)
+    else:
+        new_cell, new_shape, new_grid_start, shifts = None, None, None, None
+
+    if trim_map:
+        grid_start_for_map = new_grid_start
+        shape_for_map = new_shape
+    else:
+        grid_start_for_map = grid_start
+        shape_for_map = None
+        
+    if normalize_map and mask is not None:
+        cutoff = 0.5
+        if "DELFWT" in hkldata.df:
+            logger.writeln("Normalized Fo-Fc map requested.")
+            delfwt_map = hkldata.fft_map("DELFWT", grid_size=mask.shape)
+            masked = delfwt_map.array[mask.array>cutoff]
+            logger.writeln("   Whole volume: {} voxels".format(delfwt_map.point_count))
+            logger.writeln("  Masked volume: {} voxels (>{})".format(masked.size, cutoff))
+            global_mean = numpy.average(delfwt_map)
+            global_std = numpy.std(delfwt_map)
+            logger.writeln("    Global mean: {:.3e}".format(global_mean))
+            logger.writeln("     Global std: {:.3e}".format(global_std))
+            masked_mean = numpy.average(masked)
+            masked_std = numpy.std(masked)
+            logger.writeln("    Masked mean: {:.3e}".format(masked_mean))
+            logger.writeln("     Masked std: {:.3e}".format(masked_std))
+            #logger.writeln(" If you want to scale manually: {}".format())
+            scaled = (delfwt_map - masked_mean)/masked_std
+            hkldata.df["DELFWT"] /= masked_std # it would work if masked_mean~0
+            if omit_h_electron:
+                scaled *= -1
+                filename = "{}_normalized_fofc_flipsign.mrc".format(output_prefix)
+            else:
+                filename = "{}_normalized_fofc.mrc".format(output_prefix)
+            logger.writeln("  Writing {}".format(filename))
+            utils.maps.write_ccp4_map(filename, scaled, cell=hkldata.cell,
+                                      grid_start=grid_start_for_map, grid_shape=shape_for_map)
+
+        # Write Fo map as well
+        if "FWT" in hkldata.df:
+            fwt_map = hkldata.fft_map("FWT", grid_size=mask.shape)
+            masked = fwt_map.array[mask.array>cutoff]
+            masked_mean = numpy.average(masked)
+            masked_std = numpy.std(masked)
+            scaled = (fwt_map - masked_mean)/masked_std # does not make much sense for Fo map though
+            hkldata.df["FWT"] /= masked_std # it would work if masked_mean~0
+            filename = "{}_normalized_fo.mrc".format(output_prefix)
+            logger.writeln("  Writing {}".format(filename))
+            utils.maps.write_ccp4_map(filename, scaled, cell=hkldata.cell,
+                                      grid_start=grid_start_for_map, grid_shape=shape_for_map)
+
+    if trim_mtz and shifts is not None:
+        hkldata2 = utils.hkl.HklData(new_cell, hkldata.sg, df=None)
+        d_min = hkldata.d_min_max()[0]
+        for lab in map_labs + ["FP", "FC"]:
+            if lab not in hkldata.df: continue
+            gr = hkldata.fft_map(lab, grid_size=mask.shape)
+            gr = gemmi.FloatGrid(gr.get_subarray(new_grid_start, new_shape),
+                                 new_cell, hkldata.sg)
+            if hkldata2.df is None:
+                ad = gemmi.transform_map_to_f_phi(gr).prepare_asu_data(dmin=d_min)
+                hkldata2.merge_asu_data(ad, lab)
+            else:
+                hkldata2.df[lab] = gemmi.transform_map_to_f_phi(gr).get_value_by_hkl(hkldata2.miller_array())
+            hkldata2.translate(lab, -shifts)
+        hkldata = hkldata2
+        
+    dump_to_mtz(hkldata, map_labs, "{}.mtz".format(output_prefix))
+    if stats_str: open("{}_Fstats.log".format(output_prefix), "w").write(stats_str)
+# write_files()
+
+def write_coot_script(py_out, model_file, mtz_file, contour_fo=1.2, contour_fofc=3.0, ncs_ops=None):
+    with open(py_out, "w") as ofs:
+        ofs.write('imol = read_pdb("{}")\n'.format(model_file)) # TODO safer
+        ofs.write('imol_fo = make_and_draw_map("{}", "FWT", "PHWT", "", 0, 0)\n'.format(mtz_file))
+        ofs.write('imol_fofc = make_and_draw_map("{}", "DELFWT", "PHDELWT", "", 0, 1)\n'.format(mtz_file))
+        if contour_fo is not None:
+            ofs.write('set_contour_level_absolute(imol_fo, {:.1f})\n'.format(contour_fo))
+        if contour_fofc is not None:
+            ofs.write('set_contour_level_absolute(imol_fofc, {:.1f})\n'.format(contour_fofc))
+        if ncs_ops is not None:
+            for op in ncs_ops:
+                if op.given: continue
+                c, resid = utils.symmetry.find_center_of_origin(op.tr.mat, op.tr.vec)
+                if resid.length() > 1e-6: continue # coot does not support translation..
+                v = [y for x in op.tr.mat.tolist() for y in x] + c.tolist()
+                ofs.write("add_molecular_symmetry(imol, {})\n".format(",".join(str(x) for x in v)))
+# write_coot_script()
+
+def main(args):
+    if not args.halfmaps and not args.map:
+        raise SystemExit("Error: give --halfmaps or --map")
+
+    if not args.halfmaps and args.B is not None:
+        raise SystemExit("Error: -B only works with half maps")
+
+    if args.half1_only:
+        if not args.halfmaps:
+            raise SystemExit("--half1_only requires half maps")
+        logger.error("--half1_only specified. Half map 2 is used only for noise estimation")
+
+    if args.normalized_map:
+        logger.writeln("DeprecationWarning: --normalized_map is now on by default. This option will be removed in the future.")
+        
+    if not args.halfmaps:
+        logger.error("Warning: using --halfmaps is strongly recommended!")
+
+    st = utils.fileio.read_structure(args.model)
+    if not args.keep_charges:
+        utils.model.remove_charge([st])
+    ncs_org = gemmi.NcsOpList(st.ncs)
+    utils.model.expand_ncs(st)
+
+    if (args.omit_proton or args.omit_h_electron) and not st[0].has_hydrogen():
+        raise SystemExit("ERROR! --omit_proton/--omit_h_electron requested, but no hydrogen atoms were found.")
+
+    if args.halfmaps:
+        maps = utils.fileio.read_halfmaps(args.halfmaps, pixel_size=args.pixel_size)
+        has_halfmaps = True
+    else:
+        maps = [utils.fileio.read_ccp4_map(args.map, pixel_size=args.pixel_size)]
+        has_halfmaps = False
+
+    grid_start = maps[0][1]
+    g = maps[0][0]
+    st.spacegroup_hm = "P1"
+    st.cell = g.unit_cell
+
+    if st[0].has_hydrogen():
+        monlib = utils.restraints.load_monomer_library(st, monomer_dir=args.monlib)
+    else:
+        monlib = None
+
+    if args.mask:
+        mask = utils.fileio.read_ccp4_map(args.mask)[0]
+        if not args.no_check_mask_with_model:
+            if not utils.maps.test_mask_with_model(mask, st):
+                raise SystemExit("\nError: Model is out of mask.\n"
+                                 "Please check your --model and --mask. You can disable this test with --no_check_mask_with_model.")
+    elif args.mask_radius:
+        mask = utils.maps.mask_from_model(st, args.mask_radius, grid=g)
+        utils.maps.write_ccp4_map("mask_from_model.ccp4", mask)
+    else:
+        mask = None
+        logger.writeln("Warning: Mask is needed for map normalization. Use --mask or --mask_radius if you want normalized map.")
+
+    hkldata, map_labs, stats_str = calc_fofc(st, args.resolution, maps, mask=mask, monlib=monlib, B=args.B,
+                                             half1_only=args.half1_only, no_fsc_weights=args.no_fsc_weights,
+                                             sharpening_b=args.sharpening_b, omit_proton=args.omit_proton,
+                                             omit_h_electron=args.omit_h_electron,
+                                             source=args.source)
+    write_files(hkldata, map_labs, grid_start, stats_str,
+                mask=mask, output_prefix=args.output_prefix,
+                trim_map=args.trim, trim_mtz=args.trim_mtz, omit_h_electron=args.omit_h_electron)
+
+    py_out = "{}_coot.py".format(args.output_prefix)
+    write_coot_script(py_out, model_file=args.model,
+                      mtz_file=args.output_prefix+".mtz",
+                      contour_fo=None if mask is None else 1.2,
+                      contour_fofc=None if mask is None else 3.0,
+                      ncs_ops=ncs_org)
+    logger.writeln("\nOpen model and diffmap.mtz with COOT:")
+    logger.writeln("coot --script " + py_out)
+    if mask is not None:
+        logger.writeln("\nWant to list Fo-Fc map peaks? Try:")
+        logger.writeln("servalcat util map_peaks --map {}_normalized_fofc.mrc --model {} --abs_level 4.0".format(args.output_prefix, args.model))
+# main()
+
+if __name__ == "__main__":
+    import sys
+    args = parse_args(sys.argv[1:])
+    main(args)