import IMP
import IMP.pmi
import IMP.pmi.macros

is_mpi=True

model=IMP.Model()

mc=IMP.pmi.macros.AnalysisReplicaExchange0(model,
                                        stat_file_name_suffix="stat",
                                        global_output_directory="./post-EM",
                                        rmf_dir="rmfs/")

feature_list=["ISDCrossLinkMS_Distance_intrarb",
              "ISDCrossLinkMS_Distance_interrb",
              "ISDCrossLinkMS_Data_Score",
              "GaussianEMRestraint_None",
              "SimplifiedModel_Linker_Score_None",
              "ISDCrossLinkMS_Psi_1.0_",
              "ISDCrossLinkMS_Sigma_1_"]

mc.clustering("SimplifiedModel_Total_Score_None",
              "rmf_file",
              "rmf_frame_index",
              number_of_best_scoring_models=100,
              alignment_components=["ProtA","ProtB","ProtC"],
              distance_matrix_file="distance.rawmatrix.pkl",
              feature_keys=feature_list,
              load_distance_matrix_file=False,
              is_mpi=is_mpi,
              number_of_clusters=2,
              density_custom_ranges={"Density1":["ProtA",(1,100,"ProtB"),(200,350,"ProtB")],"Density2":["ProtC"]})