プログラム２

import os
import subprocess
import json
import numpy as np
import sys

def run_command(cmd):
    ret = subprocess.run(cmd, shell=True)
    if ret.returncode != 0:
        raise RuntimeError(f"Command failed: {cmd}")

def read_pdb_coordinates(pdb_file_path):
    coords = []
    atom_lines = []
    with open(pdb_file_path, 'r') as f:
        for line in f:
            if line.startswith("ATOM") or line.startswith("HETATM"):
                x = float(line[30:38])
                y = float(line[38:46])
                z = float(line[46:54])
                coords.append([x, y, z])
                atom_lines.append(line)
    return np.array(coords), atom_lines

def write_pdb_coordinates(atom_lines, coords, output_path):
    with open(output_path, 'w') as f:
        for line, c in zip(atom_lines, coords):
            new_line = (line[:30] +
                        f"{c[0]:8.3f}{c[1]:8.3f}{c[2]:8.3f}" +
                        line[54:])
            f.write(new_line)

def create_transformation(pixel_size_x, pixel_size_y, pixel_size_z, origin_shift):
    scale_x = 1.0 / pixel_size_x
    scale_y = 1.0 / pixel_size_y
    scale_z = 1.0 / pixel_size_z
    S = np.eye(4)
    S[0,0] = scale_x
    S[1,1] = scale_y
    S[2,2] = scale_z
    
    T = np.eye(4)
    T[0,3] = origin_shift[0]
    T[1,3] = origin_shift[1]
    T[2,3] = origin_shift[2]
    
    final_transform = S @ T
    return final_transform

def apply_transformation(coords, matrix):
    N = coords.shape[0]
    hom_coords = np.ones((N,4))
    hom_coords[:,:3] = coords
    transformed = hom_coords @ matrix
    return transformed[:,:3]

def invert_transformation(matrix):
    R = matrix[:3,:3]
    t = matrix[:3,3]
    R_inv = R.T
    t_inv = -R_inv @ t
    M_inv = np.eye(4)
    M_inv[:3,:3] = R_inv
    M_inv[:3,3] = t_inv
    return M_inv

def forward_procedure(
    mtz_file, pdb_file, phenix_path, output_map, output_transformed_pdb, 
    pixel_size_x, pixel_size_y, pixel_size_z, origin_shift, log_json
):
    # Step 1: Convert MTZ to map using phenix.mtz2map
    cmd_mtz2map = f"{phenix_path}/phenix.mtz2map {mtz_file}"
    run_command(cmd_mtz2map)
    
    # Assume the output from mtz2map is something like {basename}.map
    # If mtz_file is input.mtz, phenix.mtz2map typically generates something like input.map
    base = os.path.splitext(mtz_file)[0]
    input_map = f"{base}.map"
    
    # Step 2: Resample map to uniform pixel size using phenix.map_resample
    cmd_resample = f"{phenix_path}/phenix.map_resample {input_map} {output_map} pixel_size=1.0"
    run_command(cmd_resample)
    
    # The user must record pixel sizes and origin shift used or determined.
    # Here we assume pixel_size_x, pixel_size_y, pixel_size_z, and origin_shift are known.
    # In practice, one would parse phenix.map_resample output for these values.
    # We record them so we can invert later.
    
    # Step 3: Read original pdb coordinates
    original_coords, atom_lines = read_pdb_coordinates(pdb_file)
    
    # Step 4: Create forward transform
    T = create_transformation(pixel_size_x, pixel_size_y, pixel_size_z, origin_shift)
    
    # Step 5: Apply transformation to PDB coordinates
    transformed_coords = apply_transformation(original_coords, T)
    write_pdb_coordinates(atom_lines, transformed_coords, output_transformed_pdb)
    
    # Step 6: Save transformation data to JSON for later reversal
    data = {
        "mtz_file": mtz_file,
        "pdb_file": pdb_file,
        "phenix_path": phenix_path,
        "input_map": input_map,
        "resampled_map": output_map,
        "pixel_size_x": pixel_size_x,
        "pixel_size_y": pixel_size_y,
        "pixel_size_z": pixel_size_z,
        "origin_shift": origin_shift,
        "transform_matrix": T.tolist(),
        "transformed_pdb": output_transformed_pdb
    }
    with open(log_json, 'w') as f:
        json.dump(data, f, indent=4)

def backward_procedure(
    built_in_uniform_space_pdb, log_json, output_restored_pdb
):
    # Step 1: Load transformation data
    with open(log_json, 'r') as f:
        data = json.load(f)
    T_list = data["transform_matrix"]
    T = np.array(T_list)
    
    # Step 2: Invert transform
    T_inv = invert_transformation(T)
    
    # Step 3: Read newly built model in uniform space
    new_coords, new_atom_lines = read_pdb_coordinates(built_in_uniform_space_pdb)
    
    # Step 4: Apply inverse transform
    restored_coords = apply_transformation(new_coords, T_inv)
    write_pdb_coordinates(new_atom_lines, restored_coords, output_restored_pdb)


if __name__ == "__main__":
    # Example usage:
    # Forward:
    # python script.py forward input.mtz original.pdb /path/to/phenix output_resampled.map transformed.pdb 1.2 1.2 1.5 "-5.0 3.0 -2.0" transform_log.json
    #
    # Backward:
    # python script.py backward built_in_uniform_space.pdb transform_log.json restored_original_frame.pdb
    
    mode = sys.argv[1]
    if mode == "forward":
        mtz_file = sys.argv[2]
        pdb_file = sys.argv[3]
        phenix_path = sys.argv[4]
        output_map = sys.argv[5]
        output_transformed_pdb = sys.argv[6]
        pixel_size_x = float(sys.argv[7])
        pixel_size_y = float(sys.argv[8])
        pixel_size_z = float(sys.argv[9])
        origin_shift_str = sys.argv[10].strip()
        origin_shift = [float(x) for x in origin_shift_str.split()]
        log_json = sys.argv[11]
        forward_procedure(mtz_file, pdb_file, phenix_path, output_map, output_transformed_pdb,
                          pixel_size_x, pixel_size_y, pixel_size_z, origin_shift, log_json)
    elif mode == "backward":
        built_in_uniform_space_pdb = sys.argv[2]
        log_json = sys.argv[3]
        output_restored_pdb = sys.argv[4]
        backward_procedure(built_in_uniform_space_pdb, log_json, output_restored_pdb)
    else:
        raise ValueError("Mode must be either 'forward' or 'backward'.")