From 4e507130a59b8f15c8e0fa2997317ce03c807d17 Mon Sep 17 00:00:00 2001 From: Wong Wun-Kwint Date: Mon, 4 Aug 2025 17:09:17 +1000 Subject: [PATCH] Cprofiling the control.py --- riser/control.py | 279 +++++++++++++++++++++++++---------------------- riser/model.py | 10 +- riser/riser.py | 37 +++---- 3 files changed, 177 insertions(+), 149 deletions(-) diff --git a/riser/control.py b/riser/control.py index 99f7227..a9fec7b 100644 --- a/riser/control.py +++ b/riser/control.py @@ -1,153 +1,174 @@ import time +import cProfile +import pstats +import io +import pod5 -class SequencerControl(): - def __init__(self, client, models, processor, logger, out_file): - self.client = client + +class SequencerControl: + def __init__(self, models, processor, logger, out_file): + # removed as to simulate online self.models = models self.proc = processor self.logger = logger self.out_filename = out_file - def target(self, mode, duration_h, threshold, unblock_duration=0.1): - self.client.send_warning( - 'The sequencing run is being controlled by RISER, reads that are ' - 'not in the target class will be ejected from the pore.') - - with open(f'{self.out_filename}.csv', 'a') as out_file: + def target(self, mode, threshold, pod5_file_path, limit=None): + if limit is None: + limit = 2000 # Default limit if not provided + # No live warning or client checks + with open(f"{self.out_filename}.csv", "a") as out_file, pod5.Reader( + pod5_file_path + ) as reader: self._write_header(out_file) - run_start = time.monotonic() - progress_time = run_start + 60 - duration_s = self._hours_to_seconds(duration_h) n_assessed = 0 n_rejected = 0 n_accepted = 0 polyA_cache = {} - while self.client.is_running() and time.monotonic() < run_start + duration_s: - # Get batch of reads to process - batch_start = time.monotonic() - reads_to_reject = [] - reads_to_accept = [] - reads_unclassified = [] - for channel, read in self.client.get_read_batch(): - # Preprocess the signal - signal = self.client.get_raw_signal(read) - - # Attempt to trim the polyA tail - signal, is_polyA_trimmed = self.proc.trim_polyA(signal, read.id, polyA_cache) - - # If we haven't found the polyA yet - if not is_polyA_trimmed: - - # Use a fixed trim length if enough time has passed - if self.proc.should_trim_fixed_length(signal): - signal = self.proc.trim_polyA_fixed_length(signal) - - # Show max input length to network - signal = signal[:self.proc.get_max_length()] - - # Otherwise, try again the next time we see this read - else: - continue - - # If we have trimmed the polyA - else: - # Make sure signal is long enough to be assessed - if len(signal) < self.proc.get_min_length(): - continue - - # Trim signal if it is too long - if len(signal) > self.proc.get_max_length(): - signal = signal[:self.proc.get_max_length()] - - # Normalise - signal = self.proc.mad_normalise(signal) - - # Classify the RNA class to which the read belongs - p_on_targets = [] - p_off_targets = [] - for model in self.models: - p_off_target, p_on_target = model.classify(signal) - p_off_targets.append(p_off_target) - p_on_targets.append(p_on_target) - n_assessed += 1 - - # Decide what to do with this read - if any(p > threshold for p in p_on_targets): - decision = "accept" if mode == "enrich" else "reject" - elif all(p > threshold for p in p_off_targets): - decision = "accept" if mode == "deplete" else "reject" - elif self.proc.is_max_length(signal): - decision = "no_decision" + reads_to_reject = [] # For logging only (no actual rejection) + reads_to_accept = [] + reads_unclassified = [] + + # cProfile to see how long it takes + pr = cProfile.Profile() + pr.enable() + + # Iterate over all reads in POD5 + batch_start = time.monotonic() + for read_record in reader.reads(): + if n_assessed >= limit: + break # Stop after reaching the limit + # Preprocess the signal for pod5 + signal = read_record.signal + read_id = str(read_record.read_id) # UUID strings to plain strings + + + # Attempt to trim the polyA tail + signal, is_polyA_trimmed = self.proc.trim_polyA( + signal, read_id, polyA_cache + ) + + # If we haven't found the polyA yet + if not is_polyA_trimmed: + + # Use a fixed trim length if enough time has passed + if self.proc.should_trim_fixed_length(signal): + signal = self.proc.trim_polyA_fixed_length(signal) + + # Show max input length to network + signal = signal[: self.proc.get_max_length()] + + # Otherwise, skip (no "try again" in offline mode) else: - decision = "try_again" - - # Process the read decision - if decision == "accept": - reads_to_accept.append((channel, self._get_read_id(read))) - elif decision == "reject": - reads_to_reject.append((channel, self._get_read_id(read))) - elif decision == "no_decision": - reads_unclassified.append((channel, self._get_read_id(read))) - self._write(out_file, batch_start, channel, read.id, - len(signal), self.models, p_on_targets, - threshold, mode, decision) - - # Clear the polyA cache every 1000 reads - if len(polyA_cache) >= 1000: - polyA_cache = {} - - # Send reject requests - self.client.reject_reads(reads_to_reject, unblock_duration) - n_rejected += len(reads_to_reject) - - # Don't need to reassess the reads that were rejected, accepted - # or couldn't be classified after the maximum input length - done = reads_to_reject + reads_to_accept + reads_unclassified - self.client.finish_processing_reads(done) - n_accepted += len(reads_to_accept) - - # Log progress each minute - if batch_start > progress_time: - self.logger.info(f"In the last minute {n_assessed} signals " - f"were assessed, {n_accepted} were " - f"accepted and {n_rejected} were rejected") - n_assessed = 0 - n_rejected = 0 - n_accepted = 0 - progress_time = batch_start + 60 - else: - self.client.send_warning('RISER has stopped running.') - if not self.client.is_running(): - self.logger.info('Client has stopped.') - if time.monotonic() > run_start + duration_s: - self.logger.info(f'RISER has timed out after {duration_h} ' - 'hours as requested.') + continue + + # If we have trimmed the polyA + else: + + # Make sure signal is long enough to be assessed + if len(signal) < self.proc.get_min_length(): + continue + + # Trim signal if it is too long + if len(signal) > self.proc.get_max_length(): + signal = signal[: self.proc.get_max_length()] + + # Normalise + signal = self.proc.mad_normalise(signal) + + # Classify the RNA class to which the read belongs + p_on_targets = [] + p_off_targets = [] + for model in self.models: + p_off_target, p_on_target = model.classify(signal) + p_off_targets.append(p_off_target) + p_on_targets.append(p_on_target) + n_assessed += 1 + + # Decide what to do with this read + if any(p > threshold for p in p_on_targets): + decision = "accept" if mode == "enrich" else "reject" + elif all(p > threshold for p in p_off_targets): + decision = "accept" if mode == "deplete" else "reject" + elif self.proc.is_max_length(signal): + decision = "no_decision" + else: + decision = "try_again" + + # Process the read decision for logging + if decision == "accept": + reads_to_accept.append((channel, read_id)) + elif decision == "reject": + reads_to_reject.append((channel, read_id)) + elif decision == "no_decision": + reads_unclassified.append((channel, read_id)) + self._write( + out_file, + batch_start, + channel, + read_id, + len(signal), + self.models, + p_on_targets, + threshold, + mode, + decision, + ) + + # Clear the polyA cache every 1000 reads + if len(polyA_cache) >= 1000: + polyA_cache = {} + + + + # End the profiling + pr.disable() + s = io.StringIO() + ps = pstats.Stats(pr, stream=s).sort_stats("cumtime") + ps.print_stats() + self.logger.info("cProfile Output:\n" + s.getvalue()) + with open(f"{self.out_filename}_profile.txt", "w") as prof_file: + prof_file.write(s.getvalue()) # Save to file + + n_rejected = len(reads_to_reject) + n_accepted = len(reads_to_accept) + self.logger.info( + f"Total: {n_assessed} signals assessed, {n_accepted} accepted, {n_rejected} rejected" + ) def start(self): - self.client.start_streaming_reads() - self.logger.info('Live read stream started.') + self.logger.info("POD5 processing started.") def finish(self): - self.client.reset() - self.logger.info('Client reset and live read stream ended.') + self.logger.info("POD5 processing ended.") def _hours_to_seconds(self, hours): return hours * 60 * 60 def _get_read_id(self, read): - # Support for minknow-api <= v5.* - if hasattr(read, "number"): - return read.number - # Support for minknow-api >= v6.* - else: - return read.id + return str(read.read_id) def _write_header(self, csv_file): - csv_file.write('batch_start,read_id,channel,sig_length,models,prob_targets,threshold,mode,decision\n') - - def _write(self, csv_file, batch_start, channel, read, sig_length, - models, p_on_targets, threshold, mode, decision): - csv_file.write(f'{batch_start:.0f},{read},{channel},{sig_length},' - f'{";".join([m.target for m in models])},' - f'{";".join([str(p.item()) for p in p_on_targets])},' - f'{threshold},{mode},{decision}\n') + csv_file.write( + "batch_start,read_id,channel,sig_length,models,prob_targets,threshold,mode,decision\n" + ) + + def _write( + self, + csv_file, + batch_start, + channel, + read, + sig_length, + models, + p_on_targets, + threshold, + mode, + decision, + ): + csv_file.write( + f"{batch_start:.0f},{read},{channel},{sig_length}," + f'{";".join([m.target for m in models])},' + f'{";".join([str(p.item()) for p in p_on_targets])},' + f"{threshold},{mode},{decision}\n" + ) diff --git a/riser/model.py b/riser/model.py index 5dc1361..56b2099 100644 --- a/riser/model.py +++ b/riser/model.py @@ -16,7 +16,8 @@ def __init__(self, state, config, logger, target): # Build CNN for testing self.model = ConvNet(config.cnn).to(self.device) - self.model.load_state_dict(torch.load(state)) + # self.model.load_state_dict(torch.load(state)) + self.model.load_state_dict(torch.load(state, map_location=self.device)) self.model.eval() def classify(self, signal): @@ -30,3 +31,10 @@ def classify(self, signal): def _get_device(self): device = 'cuda' if torch.cuda.is_available() else 'cpu' return torch.device(device) + # def _get_device(self): + # # if torch.cuda.is_available(): + # # return torch.device('cuda') + # # elif torch.backends.mps.is_available(): + # # return torch.device('mps') # For Apple Silicon GPUs like M2 + # # else: + # return torch.device('cpu') diff --git a/riser/riser.py b/riser/riser.py index 77a5675..5a805b8 100644 --- a/riser/riser.py +++ b/riser/riser.py @@ -9,7 +9,6 @@ import attridict import yaml -from client import Client from model import Model from control import SequencerControl from preprocess import SignalProcessor, Kit @@ -18,12 +17,12 @@ DT_FORMAT = '%Y-%m-%dT%H:%M:%S' -def get_config(filepath): +def get_config(filepath): with open(filepath) as config_file: - return attridict(yaml.load(config_file, Loader=yaml.Loader)) + return attridict(yaml.load(config_file, Loader=yaml.Loader)) # change to dictionary style -def get_pore_version(kit): +def get_pore_version(kit): # get pore version from sequencing kit if kit == "RNA002": return "R9.4.1" elif kit == "RNA004": @@ -32,12 +31,12 @@ def get_pore_version(kit): raise Exception(f"Invalid kit {kit}") -def get_models(targets, logger, kit): +def get_models(targets, logger, kit): # load model per RNA class pore = get_pore_version(kit) models = [] for target in targets: - config = get_config(f"model/{target}_config_{kit}_{pore}.yaml") - model_file = f"model/{target}_model_{kit}_{pore}.pth" + config = get_config(f"model/{target}_config_{kit}_{pore}.yaml") # settings for the model + model_file = f"model/{target}_model_{kit}_{pore}.pth" # actual trained model models.append(Model(model_file, config, logger, target)) return models @@ -94,13 +93,6 @@ def main(): help='Whether to enrich or deplete the target class(es).' ' (required)', required=True) - parser.add_argument('-d', '--duration', - dest='duration_h', - type=float, - help='Length of time (in hours) to run RISER for. ' - 'This should be the same as the MinKNOW run ' - 'length. (required)', - required=True) parser.add_argument('-k', '--kit', choices=['RNA002', 'RNA004'], help='Sequencing kit. (required)', @@ -110,24 +102,31 @@ def main(): type=probability, help='Probability threshold for classifier [0,1] ' '(default: %(default)s)') + parser.add_argument('-f', '--pod5_file', + required=True, + help='Path to POD5 file.') + parser.add_argument('-l', '--limit', + type=int, + default=None, + help='Limit number of reads to process.') args = parser.parse_args() # Local testing # args = SimpleNamespace() # args.target = ['mRNA', 'mtRNA'] # args.mode = 'deplete' - # args.duration_h = 0.05 # args.kit = "RNA002" # args.prob_threshold = 0.9 + # args.pod5_file = "path/to/file.pod5" + # args.limit = 2000 # Set up out_file = f'riser_{get_datetime_now()}' logger = setup_logging(out_file) - client = Client(logger) models = get_models(args.target, logger, args.kit) kit = Kit.create_from_version(args.kit) processor = SignalProcessor(kit) - control = SequencerControl(client, models, processor, logger, out_file) + control = SequencerControl(models, processor, logger, out_file) # HAVE REMOVE CLIENT # Log CL args logger.info(f'Usage: {" ".join(sys.argv)}') @@ -140,9 +139,9 @@ def main(): # Run analysis control.start() - control.target(args.mode, args.duration_h, args.prob_threshold) + control.target(args.mode, args.prob_threshold, args.pod5_file, args.limit) control.finish() if __name__ == "__main__": - main() + main() \ No newline at end of file