postprocess_server.py

import h5py
import logging
import json
import time
import numpy as np
import threading
import os
import subprocess
import re
import hdf5plugin
import zmq
import argparse
import shutil
from string import digits
from glob import glob

from utility import *
from xds_ctrl import *
from dials_ctrl import *

VERSION = "for JF_GUI/v2025.09.xx or later"
V_DATE = "2025.09.07"

ROOT_DATA_SAVED = "/data/noether/jem2100plus/" # for local-test on another server
XDS_TEMPLATE = '[your_file_path]/XDS-TEMPLATE.INP'
XDS_EXE = 'xds_par'
XSCALE_EXE = 'xscale_par
SHELXT_EXE = 'shelxt'
PORT = 3467

class PostprocessLauncher:
    def __init__(self, port_number = PORT):
        self.port_number = port_number
        self.context = zmq.Context()
        self.socket = self.context.socket(zmq.REP)
        self.socket.bind(f"tcp://*:{port_number}")
        self.root_data_directory = ROOT_DATA_SAVED
        self.results = None
        self.error_retry = 5
        self.xtallist = []
        self.list_reprocess = []
        self.parameter = {}
        self.interrupt_phasing = threading.Event()
        self.ccdc_hits = []

    def update_processinfo(self, info_d, inherit=False):
        if inherit:
            self.results = info_d
        if not info_d or len(info_d) == 1: return
    
        if 'position' in info_d and 'gui_id' in info_d and not 'dataid' in info_d:
            if next((item for item in self.xtallist if item.get('gui_id') == info_d["gui_id"]), None) is None:
                self.xtallist.append(info_d)
                logging.info(f"Position item is added.: {info_d}")
            return

        prev_xtalid = next((i for i, item in enumerate(self.xtallist) if item.get('dataid') == info_d["dataid"]), None)
        if prev_xtalid is not None:
            logging.info(f"Item {info_d['dataid']} is updated")
            self.xtallist[prev_xtalid].update(info_d)
        else:
            logging.info(f"Item {info_d['dataid']} is added.")
            self.xtallist.append(info_d)
        with open(args.path_process + '/process_result.jsonl', 'w') as f:
            [f.write(json.dumps(i, cls=CustomJSONEncoder) + "\n") for i in self.xtallist]
            logging.info(f"{args.path_process + '/process_result.jsonl'} is updated.")
            
    def pick_lastimages(self):
        filepaths = [item.get('filepath') for item in self.xtallist if item.get('filepath') is not None]
        hdfs = sorted(glob(os.path.dirname(filepaths[-1]) + '/*master.h5'))[::-1]
        subimages = None
        for hdf in hdfs:
            packed_data = pick_subimages(hdf)
            if packed_data is not None:
                logging.info(f'Reload images from {hdf}')
                break
        return packed_data
    
    def run_phasing(self, event, thread_watch, timeout, composition, processor='x'):
        logging.info('Waiting until HKL file is ready...')
        while not event.wait(timeout):
            if not thread_watch.is_alive():
                if processor=='x' and os.path.isfile(args.path_process + '/XSCALE/xds_shelx.hkl'):
                    logging.info('Launching phasing routine...')
                    run_shelxt(args.path_process + '/XSCALE/SHELX',
                               args.path_process + '/XSCALE/xds_shelx.hkl',
                               'crystal', SHELXT_EXE, args.quiet)
                    return
                elif processor=='d' and os.path.isfile(args.path_process + '/DIALS/xia2/SHELX/dials.hkl'):
                    logging.info('Launching phasing routine...')
                    run_shelxt(args.path_process + '/DIALS/xia2/SHELX',
                               args.path_process + '/DIALS/xia2/SHELX/dials.hkl',
                               'dials', SHELXT_EXE, args.quiet)
                    return
                # else:
                #     logging.error(f'Undefined processor ID: {processor}!')
                #     return
        logging.warning('Waiting for phasing has timed-out.')
        event.clear()
        
    def run(self):
        logging.info("Server started, waiting for launching postprocess requests...")
        logging.info(f"Args: {args}")
        ccdcscreener = CCDCscreeningInquiry()
        
        try:
            with open(args.path_process + '/process_result.jsonl', 'r') as f:
                prev_xtallist = [json.loads(line) for line in f]
                self.xtallist = [d for d in prev_xtallist if not (len(d) == 1 and 'filename' in d)]
                logging.info(f"Loaded the session-data: N={len(self.xtallist)}")
                self.ccdc_hits = ccdcscreener.run(self.xtallist)
                if len(prev_xtallist[-1]) == 1 and 'filename' in prev_xtallist[-1]: self.xtallist.append(prev_xtallist[-1])
        except FileNotFoundError:
            logging.warning(f"No session-data loaded.")
        
        while True:
            ready_postprocess = False
            ready_reprocess = False
            xds_reprocess = False
            dials_reprocess = False
            ready_merge = False
            xds_merge = False
            dials_merge = False
            try:
                message_raw = self.socket.recv_string()
                # start usual post-processing with a pair of [hdf-filepath, gui-id]
                if 'Launching the postprocess...' in message_raw:
                    print(message_raw)
                    filename = self.root_data_directory + message_raw.split()[-3]
                    gui_id = int(message_raw.split()[-1])
                    if not os.path.exists(filename):
                        logging.error(f"{filename} is not found!!")
                        self.socket.send_string('Saved data not found')
                    else:
                        self.socket.send_string(f'{filename} will be processed soon.')
                        ready_postprocess = True
                # send back the single post-processing result
                elif args.feedback and 'Results being inquired...' in message_raw:
                    print(message_raw)
                    if self.results is None:
                        self.socket.send_string("In processing...")
                    else:
                        logging.info("Sending results to GUI...")
                        self.socket.send_string(json.dumps(self.results, cls=CustomJSONEncoder))
                        self.results = None
                # load the session data from the file, and hold it on memory
                elif args.feedback and 'Session-metadata being inquired...' in message_raw:
                    print(message_raw)
                    session_file = args.path_process + '/process_result.jsonl'
                    if not os.path.exists(session_file):
                        logging.info(f"Previous session-data does not exist as {session_file}")
                        self.socket.send_string('Previous session-data not found')
                        return
                    if len(self.xtallist) < 1:
                        logging.info(f"Load the session-file: {session_file}")
                    else:
                        logging.info(f"Reload the session-file: {session_file}")
                    with open(session_file, 'r') as f:
                        self.xtallist = [json.loads(line) for line in f]
                    self.ccdc_hits = ccdcscreener.run(self.xtallist, timeout_ms=1000)
                    logging.info("Sending session-data...")
                    if 'False' in message_raw:
                        subimages = {"image_supporting": self.pick_lastimages()}
                        self.socket.send_string(json.dumps(self.xtallist+self.ccdc_hits+[subimages], cls=CustomJSONEncoder))
                    else:
                        self.socket.send_string(json.dumps(self.xtallist+self.ccdc_hits, cls=CustomJSONEncoder))
                else:
                    message = json.loads(message_raw)
                    # launch re-processing
                    if 'dataids' in message:
                        print(message)
                        self.list_reprocess = []
                        self.parameter = {}
                        # load and send back the spot-data to GUI
                        if 'nskip' in message:
                            xtals = []
                            for dataid in message['dataids']:
                                dic = [item for item in self.xtallist if item.get('dataid') == dataid][0]
                                if 'beamcenter_used' in dic:
                                    spotloader = SpotLoader(dic['filepath'], nskip=message['nskip'], beamcenter=dic['beamcenter_used'])
                                else:
                                    spotloader = SpotLoader(dic['filepath'], nskip=message['nskip'])
                                cell_axes = np.array(dic.get('cell axes'), dtype=float).reshape((3,3))
                                xtals.append({"spots_jfj": spotloader.reshape_3dreal(),
                                             "dataids": dataid,
                                             "cell axes": cell_axes})
                            logging.info("Sending spots-data...")
                            self.socket.send_string(json.dumps(xtals, cls=CustomJSONEncoder))
                        elif 'formula' in message:
                            for i in message['dataids']:
                                if os.path.exists(args.path_process + '/XDS/' + i + '/XDS_ASCII.HKL'):
                                    self.list_reprocess.append(args.path_process + '/XDS/' + i + '/XDS_ASCII.HKL')
                            if len(self.list_reprocess) > 0:
                                self.parameter = message
                                self.socket.send_string('Merge-and-solve signal was caught successfully.')
                                ready_merge = True
                                xds_merge = message['xds']
                                dials_merge = message['dials']
                            else:
                                self.socket.send_string(f'No reprocessed data was found.')
                        else:
                            for i in message['dataids']:
                                if os.path.exists(args.path_process + '/XDS/' + i + '/XDS.INP'):
                                    # self.list_reprocess.append(args.path_process + '/XDS/' + i + '/XDS.INP')
                                    self.list_reprocess.append(i)
                            if len(self.list_reprocess) > 0:
                                self.parameter = message
                                self.socket.send_string('Reprocess signal was caught successfully.')
                                ready_reprocess = True
                                xds_reprocess = message['xds']
                                dials_reprocess = message['dials']
                            else:
                                self.socket.send_string(f'No preprocessed data was found.')
                    # receive the position data and save it in a json file with others
                    else:
                        if isinstance(message, list) and args.json:
                            [self.update_processinfo(i) for i in message]
                            self.socket.send_string("Position-info added successfully")
                        else:
                            logging.error(f"Received undefined json-data: {message}")
                            logging.error(f"Otherwise, Json-writing option ('-j') is missing")
            except zmq.ZMQError as e:
                logging.error(f"Error while receiving request: {e}")
                self.error_retry -= 1
                if self.error_retry < 0: break
            except Exception as e:
                logging.error(f"Error while receiving/processing request: {e}", exc_info=True)
            #     break

            if ready_merge:
                composition = self.parameter["formula"].translate(str.maketrans("", "", digits))
                if 'x' in args.processor and xds_merge:

                    xmerge_thread = threading.Thread(target=merge_xds, 
                        args=(args.path_process + '/XDS', self.parameter["dataids"], XSCALE_EXE, args.quiet, 0, self.parameter["spacegroup"], [20, 0.7], self.update_processinfo, composition, ), daemon=True) # use averaged cell and for resolution of 0-inf.
                    xmerge_thread.start()
                    # xmerge_thread.join(timeout=10)

                    phasing_thread = threading.Thread(target=self.run_phasing,
                        args=(self.interrupt_phasing, xmerge_thread, 5, composition, 'x',), daemon=True)
                    timeout=60
                if 'd' in args.processor and dials_merge:
                    dmerge_thread = threading.Thread(target=merge_by_xia2, 
                        args=(args.path_process + '/DIALS', self.parameter["dataids"], self.parameter["spacegroup"], self.update_processinfo, composition, args.quiet, ), daemon=True)
                    dmerge_thread.start()

                    phasing_thread = threading.Thread(target=self.run_phasing,
                        args=(self.interrupt_phasing, dmerge_thread, 5, composition, 'd',), daemon=True)
                    timeout=180
                    
                phasing_thread.start()
                timeout_phasing = threading.Timer(timeout, self.interrupt_phasing.set)
                timeout_phasing.start()

            if ready_reprocess:
                if 'x' in args.processor and xds_reprocess:
                    for dataid in self.list_reprocess:
                        filename = args.path_process + '/XDS/' + dataid + '/XDS.INP'
                        xds_thread = threading.Thread(target=rerun_xds, 
                                  args=(filename, XDS_EXE, args.quiet, self.parameter["cell"], self.parameter["spacegroup"], self.update_processinfo, ), daemon=True)
                        xds_thread.start()
                if 'd' in args.processor and dials_reprocess:
                    for dataid in self.list_reprocess:
                        dials_thread = threading.Thread(target=rerun_dials, 
                                  args=(args.path_process + '/DIALS/' + dataid, args.quiet, self.parameter["cell"], self.parameter["spacegroup"], self.update_processinfo, ), daemon=True)
                        dials_thread.start()
                    
            if ready_postprocess:
                with h5py.File(filename, 'r') as f:
                    beamcenter = f['entry/instrument/detector/beam_center_x'][()], f['entry/instrument/detector/beam_center_y'][()]
                    if beamcenter[0]*beamcenter[1] == 1 or args.refinecenter:
                        beamcenter = refine_beamcenter(f)

                dataid = re.sub(".*/([0-9]{3})_.*_([0-9]{4})_master.h5","\\1-\\2", filename)
                logging.info(f'Subdirname: {os.path.dirname(filename)}/XDS/{dataid}')

                process_dir = args.path_process
                if process_dir == '.':
                    process_dir = os.path.dirname(filename)

                if 'x' in args.processor:
                    myxds = XDSparams(XDS_TEMPLATE, XDS_EXE, center_margin=args.center_mask)
                    xds_thread = threading.Thread(target=myxds.run_xds, 
                        args=(filename, process_dir + '/XDS/' + dataid,
                                    beamcenter, args.quiet, 
                                    args.exoscillation, self.update_processinfo, gui_id, ), 
                            daemon=True)
                    xds_thread.start()

                if 'd' in args.processor:
                    mydials = DIALSparams(filename, process_dir + '/DIALS/' + dataid, 
                                  beamcenter=beamcenter)
                    dials_thread = threading.Thread(target=mydials.run_dials, 
                        args=(server.update_processinfo, args.quiet, gui_id, ), daemon=True)
                    dials_thread.start()                    
    
    def stop(self):
        self.running = False
        logging.info("Stopping server...")

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("-b", "--feedback", action="store_true", help="send post-process results to GUI")
    parser.add_argument("-c", "--center_mask", type=int, default=10, help="centre mask radius [px] in XDS.INP (10). deactivate with 0.")    
    parser.add_argument("-d", "--path_process", type=str, default='.', help="root directory path for data-processing (.). file-writing permission is necessary.")
    parser.add_argument("-j", "--json", action="store_true", help="write a summary of postprocess as a JSON'L' file (process_result.jsonl)")
    parser.add_argument("-o", "--exoscillation", action="store_true", help="use measured oscillation value for postprocess")
    # parser.add_argument("-op", "--opticalcenter", action="store_true", help="calibtate the beam direction")
    parser.add_argument("-q", "--quiet", action="store_true", help="suppress outputs of external programs")
    parser.add_argument("-r", "--refinecenter", action="store_true", help="force post-refine beamcenter position")
    parser.add_argument("-v", "--version", action="store_true", help="display version information")
    # parser.add_argument("-f", "--formula", type=str, default='C2H5NO2', help="chemical formula for ab-initio phasing with shelxt/d")
    parser.add_argument("-p", "--processor", type=str, default='x', help="enable post-processing. 'x' for XDS, 'd' for dials and 'xd' for both")
    parser.add_argument("-f", "--filepath", type=str, default='', help="offline analysis. identify the full-filepath to be processed")

    args = parser.parse_args()

    print(f"Postprocessing-server for diffraction datasets recorded by Jungfrau / Jungfrau-joch: ver. {V_DATE} {VERSION}")

    if args.version:
        print('''
            **Detailed information of authors, years, project name, Github URL, license, contact address, etc.**
            Postprocess-server for Electron Diffraction with JUNGFRAU (2024-)
            https://github.com/epoc-ed/GUI
            EPOC Project (2024-2025, UniVie - PSI)
            https://github.com/epoc-ed
            https://epoc-ed.github.io/manual/index.html
        ''')
    
    # Initialize logger
    logger = logging.getLogger()

    logger.setLevel("INFO")

    # Create the handler for console output
    console_handler = logging.StreamHandler()

    # Apply the custom formatter to the handler
    formatter = CustomFormatter('%(asctime)s - %(levelname)s - %(message)s', datefmt='%H:%M:%S')
    console_handler.setFormatter(formatter)

    # Add the handler to the logger
    logger.addHandler(console_handler)

    if not os.access(args.path_process, os.W_OK):
        logging.warning("No file-writing permission!")
        exit()
    
    server = PostprocessLauncher()
    if args.filepath == '':
        server_thread = threading.Thread(target=server.run, daemon=True)
        server_thread.start()
        logging.info("Postprocessing-server is running in the background. You can now use the command line.")
        # The server will continue to run in the background until exit() is called
    else:
        filename = args.filepath

        try:
            with open(args.path_process + '/process_result.jsonl', 'r') as f:
                prev_xtallist = [json.loads(line) for line in f]
                server.xtallist = [d for d in prev_xtallist if not (len(d) == 1 and 'filename' in d)]
                logging.info(f"Loaded the session-data: N={len(server.xtallist)}")
                # self.ccdc_hits = ccdcscreener.run(self.xtallist)
                if len(prev_xtallist[-1]) == 1 and 'filename' in prev_xtallist[-1]: server.xtallist.append(prev_xtallist[-1])
        except FileNotFoundError:
            logging.warning(f"No session-data loaded.")
        
#         # reprocessing
#         if 'd' in args.processor:
#             rerun_dials(filename, cell=[5.8,9.5,11.5,90,90,90], spacegroup=16)
        
#         exit()
        
        # initial processing
        with h5py.File(filename, 'r') as f:
            beamcenter = f['entry/instrument/detector/beam_center_x'][()], f['entry/instrument/detector/beam_center_y'][()]
            if beamcenter[0]*beamcenter[1] == 1 or args.refinecenter:
                beamcenter = refine_beamcenter(f)

        process_dir = args.path_process
        if process_dir == '.':
            process_dir = os.path.dirname(filename)

        dataid = re.sub(".*/([0-9]{3})_.*_([0-9]{4})_master.h5","\\1-\\2", filename)
        logging.info(f'Subdirname: {process_dir}/XDS/{dataid}')
        
        gui_id_dummy = int(os.path.basename(filename)[:3])+1000
        
        if 'x' in args.processor:
            myxds = XDSparams(XDS_TEMPLATE, XDS_EXE, center_margin=args.center_mask)
            myxds.run_xds(filename, process_dir + '/XDS/' + dataid, 
                           beamcenter, args.quiet, args.exoscillation, 
                           server.update_processinfo, gui_id_dummy, True, False) # pos-output, dry-run
        if 'd' in args.processor:
            mydials = DIALSparams(filename, process_dir + '/DIALS/' + dataid, 
                                  beamcenter=beamcenter)
            mydials.run_dials(server.update_processinfo, args.quiet, gui_id_dummy, True, False) # pos-output, dry-run