Skip to content

Cprofiling the control.py #4

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
kyawzinhtun7415 wants to merge 1 commit into
base: develop
Choose a base branch
from
Open

Cprofiling the control.py #4

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
279 changes: 150 additions & 129 deletions riser/control.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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"
)
10 changes: 9 additions & 1 deletion riser/model.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand All @@ -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')
Loading