def main(): IMP.base.add_string_flag("input", "", "The RMF file to add cylinders to.") IMP.base.add_float_flag("radius", 5, "The radius of the cylinder.") IMP.base.add_float_flag("site_radius", 2, "The radius of the sites.") IMP.base.add_bool_flag("recolor_fgs", "recolor fg nup chains") IMP.base.add_bool_flag("recolor_floats", "recolor floating (diffusing) molecules") IMP.base.setup_from_argv(sys.argv, "Prettify a movie") fh = RMF.open_rmf_file(IMP.base.get_string_flag("input")) radius = IMP.base.get_float_flag("radius") print "opened", fh.get_name() cf = RMF.CylinderFactory(fh) rff = RMF.ReferenceFrameFactory(fh) tf = RMF.TypedFactory(fh) bf = RMF.BallFactory(fh) cdf = RMF.ColoredFactory(fh) ipf = RMF.IntermediateParticleFactory(fh) fg_types = [ IMP.npctransport.get_type_of_fg(i).get_string() for i in range(0, IMP.npctransport.get_number_of_types_of_fg()) ] float_types = [ IMP.npctransport.get_type_of_float(i).get_string() for i in range(0, IMP.npctransport.get_number_of_types_of_float()) ] fh.set_current_frame(RMF.ALL_FRAMES) if (IMP.base.get_bool_flag("recolor_fgs")): _recolor(fh.get_root_node(), tf, cdf, fg_types, fg_color) if (IMP.base.get_bool_flag("recolor_floats")): _recolor(fh.get_root_node(), tf, cdf, float_types[0:1], kap_color) _recolor(fh.get_root_node(), tf, cdf, float_types[1:], crap_color) _resize_sites(fh.get_root_node(), bf, IMP.base.get_float_flag("site_radius")) cylinders = _add_nodes(fh.get_root_node(), cf, cdf, tf, fg_types, radius, fg_color) for i in range(0, fh.get_number_of_frames()): print "frame", i fh.set_current_frame(i) for c in cylinders: _set_cylinder(c, cf, rff)
# This simple example makes an RMF file with several rigid copies of the # same thing from __future__ import print_function import RMF file_name = RMF._get_temporary_file_path("reference_frame.rmfz") print("file is", file_name) fh = RMF.create_rmf_file(file_name) fh.add_frame("first frame", RMF.FRAME) rh = fh.get_root_node() reference_frame_factory = RMF.ReferenceFrameFactory(fh) ball_factory = RMF.BallFactory(fh) particle_factory = RMF.ParticleFactory(fh) color_factory = RMF.ColoredFactory(fh) bond_factory = RMF.BondFactory(fh) origin = rh.add_child("origin", RMF.REPRESENTATION) pd = particle_factory.get(origin) pd.set_radius(1) pd.set_mass(2) pd.set_coordinates(RMF.Vector3(0, 0, 0)) for i in range(0, 3): c = rh.add_child(str(i), RMF.REPRESENTATION) pd = particle_factory.get(c) pd.set_radius(1) pd.set_mass(2) coords = [0, 0, 0] coords[i] = 2
def main(): IMP.add_string_flag("input_rmf", "", "The input RMF file.") IMP.add_string_flag("output_rmf", "", "The output RMF file in which to add cylinders.") IMP.add_string_flag( "ref_output", "", "reference output file from which info e.g. fg nup types can be extracted" ) IMP.add_float_flag("radius", 5, "The radius of the cylinder.") IMP.add_float_flag("site_radius", 2, "The radius of the sites.") IMP.add_bool_flag("recolor_fgs", "recolor fg nup chains") IMP.add_bool_flag("recolor_floats", "recolor floating (diffusing) molecules") IMP.add_int_flag( "smooth_n_frames", 1, "smooth by averaging over a window of specified frames (but retaining the same number of expected frames" ) IMP.add_int_flag("skip_n_frames", 1, "skip every n frames of output") IMP.setup_from_argv(sys.argv, "Prettify a movie") in_fname = IMP.get_string_flag("input_rmf") out_fname = IMP.get_string_flag("output_rmf") ref_output = IMP.get_string_flag("ref_output") radius = IMP.get_float_flag("radius") # Prepare out file with same static info as in file: in_fh = RMF.open_rmf_file_read_only(in_fname) out_fh = RMF.create_rmf_file(out_fname) print("creating file", out_fname) RMF.clone_file_info(in_fh, out_fh) RMF.clone_hierarchy(in_fh, out_fh) RMF.clone_static_frame(in_fh, out_fh) print("opened", in_fh.get_name()) cf = RMF.CylinderFactory(out_fh) rff = RMF.ReferenceFrameFactory(out_fh) tf = RMF.TypedFactory(out_fh) bf = RMF.BallFactory(out_fh) cdf = RMF.ColoredFactory(out_fh) ipf = RMF.IntermediateParticleFactory(out_fh) fg_types, kap_types, inert_types = _get_fg_and_floater_types(ref_output) # out_fh.set_current_frame(RMF.ALL_FRAMES) # Modify static information: if (IMP.get_bool_flag("recolor_floats")): _recolor(out_fh.get_root_node(), tf, cdf, kap_types, kap_color) _recolor(out_fh.get_root_node(), tf, cdf, inert_types, inert_color) _resize_sites(out_fh.get_root_node(), bf, IMP.get_float_flag("site_radius")) cylinders = [] for i, fg_type in enumerate(fg_types): color = IMP.display.get_display_color(i) rgb = [color.get_red(), color.get_green(), color.get_blue()] print("Checking fg type", fg_type) cylinders += _add_nodes(out_fh.get_root_node(), cf, cdf, tf, [fg_type], radius, rgb) # fg_color if (IMP.get_bool_flag("recolor_fgs")): # print "Recoloring",fg_type, rgb _recolor(out_fh.get_root_node(), tf, cdf, [fg_type], rgb) # fg_color # Clone and modify per-frame information: smooth_xyz_dict = {} smooth_n_frames = IMP.get_int_flag("smooth_n_frames") skip_n_frames = IMP.get_int_flag("skip_n_frames") print("Skip interval:", skip_n_frames, "frames") for f_id, f in enumerate(in_fh.get_frames()): is_write = f_id % skip_n_frames == 0 in_fh.set_current_frame(f) if not is_write: if (smooth_n_frames > 1): _smooth(in_fh.get_root_node(), tf, rff, smooth_xyz_dict, n=smooth_n_frames, is_write=False) print("skipping frame", f, f_id) continue print("cloning frame", f, f_id) out_fh.add_frame(in_fh.get_name(f), in_fh.get_type(f)) RMF.clone_loaded_frame(in_fh, out_fh) if (smooth_n_frames > 1): _smooth(out_fh.get_root_node(), tf, rff, smooth_xyz_dict, n=smooth_n_frames, is_write=True) for c in cylinders: _set_cylinder(c, cf, rff) if (IMP.get_bool_flag("recolor_fgs")): _recolor_cylinder(c, cf, cdf) DEBUG = False if DEBUG and f_id >= 5 * skip_n_frames: break
# find the name for a temporary file to use to for writing the hdf5 file tfn = RMF._get_temporary_file_path("aliases.rmf") print("File is", tfn) # open the temporary file, clearing any existing contents fh = RMF.create_rmf_file(tfn) # add a frame to the file fh.add_frame("root", RMF.FRAME) rh = fh.get_root_node() # add a bunch of "molecules" with colors mrh = rh.add_child("molecules", RMF.REPRESENTATION) pf = RMF.ParticleFactory(fh) cf = RMF.ColoredFactory(fh) red = [] green = [] blue = [] for i in range(0, 10): curph = mrh.add_child("mol" + str(i), RMF.REPRESENTATION) for j in range(0, 10): curh = curph.add_child("atom" + str(j), RMF.REPRESENTATION) p = pf.get(curh) p.set_coordinates(RMF.Vector3(3 * i, 3 * j, 0)) p.set_mass(1) p.set_radius(1) c = cf.get(curh) if j % 3 == 0: c.set_rgb_color(RMF.Vector3(1, 0, 0))