"""Stage to end after for debugging and runtime testing purposes."""

DC_INPUT = 1

TF_EXAMPLES = 2

RUN_MODEL = 3

flags.mark_flags_as_required([

'subreads_to_ccs',

'ccs_bam',

'checkpoint',

'output',

"""A central place to define options used across various stages of inference.

Attributes:

example_height: Height of examples, which depends on max_passes.

max_length: Length of window.

max_passes: Max number of subreads to include in input shown to model.

min_quality: Quality threshold to filter final reads.

min_length: Length threshold to filter final reads.

batch_size: Number of examples passed through model at once.

use_ccs_bq: Use CCS Base Quality Scores as a feature.

cpus: Number of processes to use for multiprocessing. Must be positive (for

multiprocessing) or 0 (for serial execution).

skip_windows_above: Run the model only when the avg(ccs_base_qual) of the

window is below this value.

use_saved_model: True if the given checkpoint is a saved model, false if it

is a regular checkpoint.

max_base_quality: The maximum base quality value allowed.

dc_calibration_values: QualityCalibrationValues defining values to be used

for deepconsensus quality calibration.

ccs_calibration_values: QualityCalibrationValues defining values to be used

for ccs quality calibration.

"""

max_length: int

example_height: int

max_passes: int

min_quality: int

min_length: int

batch_size: int

use_ccs_bq: bool

cpus: int

skip_windows_above: int

use_saved_model: bool

max_base_quality: int

dc_calibration_values: calibration_lib.QualityCalibrationValues

stage: str,

item: str,

before: float,

num_examples: Optional[int] = None,

num_subreads: Optional[int] = None,

num_zmws: Optional[int] = None,

) -> None:

"""Catalogue time elapsed for a given stage relative to "before"."""

after = time.time()

datum = {

'item': item,

'stage': stage,

'runtime': after - before,

'num_zmws': num_zmws,

'num_examples': num_examples,

'num_subreads': num_subreads,

}

feature_dicts: List[Tuple[str, Union[np.ndarray, int, bytes, float]]],

model_params: Union[config_dict.ConfigDict, config_dict.FrozenConfigDict],

options: InferenceOptions,

):

"""Stack values for each feature.

Args:

feature_dicts: List of feature dictionaries.

model_params: Parameters for the model.

options: Some options that apply to various stages of the inference run.

Yields:

For Dictionary of batched features.

"""

def process_feature_dicts(features):

return data_providers.process_feature_dict(

features=features, params=model_params

)

def split_list(l, batch_size):

for i in range(0, len(l), batch_size):

yield l[i : i + batch_size]

processed_feature_dicts = list(map(process_feature_dicts, feature_dicts))

for one_batch in split_list(processed_feature_dicts, options.batch_size):

examples = {}

for key in dc_constants.DC_FEATURES:

vals = [x[key] for x in one_batch]

if vals:

examples[key] = np.stack(vals)

else:

examples[key] = vals

feature_dicts: List[Tuple[str, Union[np.ndarray, int, bytes]]],

model: tf.keras.Model,

model_params: Union[config_dict.ConfigDict, config_dict.FrozenConfigDict],

options: InferenceOptions,

) -> List[stitch_utils.DCModelOutput]:

"""Runs the model on one example to get one predicted output sequence.

Args:

feature_dicts: List of feature dictionaries.

model: An initialized model that will be used to make predictions.

model_params: Parameters for the model.

options: Some options that apply to various stages of the inference run.

Returns:

A DeepConsensusInput proto containing the prediction from the model.

"""

predictions = []

for data in batch_examples(feature_dicts, model_params, options):

window_pos_arr = data['window_pos']

molecule_name_arr = data['name']

rows = tf.convert_to_tensor(data['rows'])

if options.use_saved_model:

softmax_output = model.signatures['serving_default'](rows)

softmax_output = softmax_output['output_1']

else:

softmax_output = model.predict(rows)

softmax_output = softmax_output.numpy()

ec_arr = data['ec']

np_num_passes_arr = data['np_num_passes']

rq_arr = data['rq']

rg_arr = data['rg']

y_preds = np.argmax(softmax_output, -1)

error_prob = 1 - np.max(softmax_output, axis=-1)

quality_scores = -10 * np.log10(error_prob)

if options.dc_calibration_values.enabled:

quality_scores = calibration_lib.calibrate_quality_scores(

quality_scores, options.dc_calibration_values

)

# The maximum value for base quality scores is max_base_quality.

quality_scores = np.minimum(quality_scores, options.max_base_quality)

quality_scores = np.round(quality_scores, decimals=0)

quality_scores = quality_scores.astype(dtype=np.int32)

# The minimum value for base quality scores is 0.

quality_scores = np.maximum(quality_scores, 0)

for y_pred, qs, window_pos, molecule_name, ec, np_, rq, rg in zip(

y_preds,

quality_scores,

window_pos_arr,

molecule_name_arr,

ec_arr,

np_num_passes_arr,

rq_arr,

rg_arr,

):

dc_output = stitch_utils.DCModelOutput(

window_pos=window_pos,

molecule_name=molecule_name,

ec=ec,

np_num_passes=np_,

rq=rq,

rg=rg,

)

y_pred_bases = ''.join(

np.vectorize(dc_constants.SEQ_VOCAB.__getitem__)(y_pred)

)

quality_string = utils.quality_scores_to_string(qs)

dc_output.sequence = y_pred_bases

dc_output.quality_string = quality_string

predictions.append(dc_output)

predictions: Iterable[stitch_utils.DCModelOutput],

zmw: str,

options: InferenceOptions,

outcome_counter=stitch_utils.OutcomeCounter,

) -> Optional[str]:

"""Stitches together predictions into one sequence.

Args:

predictions: Predictions from running model on examples.

zmw: Molecule name, the part that is shared among all subreads.

options: Options here are used for filtering.

outcome_counter: Keeps track of how many ZMWs end up with which outcomes.

Returns:

Fastq string for one sequence.

"""

fastq_string = stitch_utils.stitch_to_fastq(

molecule_name=zmw,

predictions=predictions,

max_length=options.max_length,

min_quality=options.min_quality,

min_length=options.min_length,

outcome_counter=outcome_counter,

)

subreads_to_ccs: str, ccs_bam: str, options: InferenceOptions

) -> Generator[Tuple[str, str, Sequence[Any], np.ndarray], None, None]:

"""Streams inputs from FASTA and BAM concurrently.

Args:

subreads_to_ccs: Path to input BAM file with subreads aligned to template

sequences.

ccs_bam: Path to the input CCS BAM with template sequences (e.g. CCS or

POA).

options: Inference options, used to initialize a DcConfig object.

Yields:

For every ZMW, (ZMW name, template sequence, list of subreads).

"""

dc_config = pre_lib.DcConfig(

max_passes=options.max_passes,

max_length=options.max_length,

use_ccs_bq=options.use_ccs_bq,

)

# Temporarily disable unused-variable.

# pylint: disable=unused-variable

proc_feeder, main_counter = pre_lib.create_proc_feeder(

subreads_to_ccs=subreads_to_ccs,

ccs_bam=ccs_bam,

dc_config=dc_config,

ins_trim=FLAGS.ins_trim,

use_ccs_smart_windows=FLAGS.use_ccs_smart_windows,

)

# pylint: enable=unused_variable

for input_data in proc_feeder():

subreads, zmw, dc_config, split, window_widths = input_data

checkpoint_path: str,

params: config_dict.ConfigDict,

options: InferenceOptions,

) -> Tuple[Optional[tf.keras.Model], Optional[config_dict.ConfigDict]]:

"""Initializes the model and gathers parameters.

Args:

checkpoint_path: Path to model checkpoint.

params: Parameter object, from flags.

options: Contains a few more parameters some of which will replace those in

the params object.

Returns:

A tuple containing an initialized model and a final parameter set.

"""

if FLAGS.end_after_stage in [DebugStage.TF_EXAMPLES, DebugStage.DC_INPUT]:

return None, None

model_utils.modify_params(

params=params,

speedy=True,

max_length=options.max_length,

is_training=False,

)

logging.info('Loading %s', checkpoint_path)

if options.use_saved_model:

model = tf.saved_model.load(checkpoint_path)

else:

model = model_utils.get_model(params)

# This loads a model saved in tf.train.Checkpoint format through the custom

# training loop code.

checkpoint = tf.train.Checkpoint(model=model)

# Note that the `print_model_summary` is necessary because we need to run a

# forward pass with the model in order for assert_existing_objects_matched

# to work as expected.

# If you don't do this, then assert_existing_objects_matched will not

# raise an error even if the wrong checkpoint is used.

# Some context here: b/148023980.

input_shape = (1, params.total_rows, params.max_length, params.num_channels)

model_utils.print_model_summary(model, input_shape)

checkpoint.restore(

checkpoint_path

).expect_partial().assert_existing_objects_matched()

logging.info('Finished initialize_model.')

one_zmw: Tuple[str, List[pre_lib.Read], pre_lib.DcConfig, np.ndarray]

) -> Tuple[List[Dict[str, Any]], Optional[Any]]:

"""Preprocess input data for one ZMW into windows of features.

This is often run from multiple processes in parallel, which creates some

constraints to keep in mind. Adjustments include returning runtime data

points instead of updating a global variable.

Args:

one_zmw: Input data for one ZMW: a tuple of (name, subreads, DcConfig).

Returns:

A list of feature dictionaries, one for each window.

collections.Counter with inference counters.

"""

zmw, subreads, dc_config, window_widths = one_zmw

dc_whole_zmw = pre_lib.subreads_to_dc_example(

subreads=subreads,

ccs_seqname=zmw,

dc_config=dc_config,

window_widths=window_widths,

)

if dc_whole_zmw is None or FLAGS.end_after_stage == DebugStage.DC_INPUT:

return ([], None)

# One feature dictionary per window/example.

feature_dicts = [x.to_features_dict() for x in dc_whole_zmw.iter_examples()]

feature_dict: Dict[str, Any], options: InferenceOptions

) -> stitch_utils.DCModelOutput:

"""Process a window by simply adopting the CCS sequence and base qualities."""

rows = feature_dict['subreads']

_, _, _, _, ccs_index, _, _ = data_providers.get_indices(

options.max_passes, options.use_ccs_bq

)

ccs = rows[ccs_index[0], :, 0]

ccs_seq = utils.encoded_sequence_to_string(ccs)

ccs_quality_scores = feature_dict['ccs_base_quality_scores']

if options.ccs_calibration_values.enabled:

ccs_quality_scores = calibration_lib.calibrate_quality_scores(

ccs_quality_scores, options.ccs_calibration_values

)

ccs_quality_scores = np.minimum(ccs_quality_scores, options.max_base_quality)

ccs_quality_scores = ccs_quality_scores.astype(dtype=np.int32)

dc_output = stitch_utils.DCModelOutput(

window_pos=feature_dict['window_pos'],

molecule_name=feature_dict['name'],

sequence=ccs_seq,

quality_string=utils.quality_scores_to_string(ccs_quality_scores),

ec=feature_dict['ec'],

np_num_passes=feature_dict['np_num_passes'],

rq=feature_dict['rq'],

rg=feature_dict['rg'],

)

inputs: Sequence[Tuple[str, str, Sequence[Any], np.ndarray]],

model: tf.keras.Model,

model_params: Union[config_dict.ConfigDict, config_dict.FrozenConfigDict],

output_writer: Union[gfile.GFile, pysam.AlignmentFile],

options: InferenceOptions,

batch_name: str,

outcome_counter: stitch_utils.OutcomeCounter,

stats_counter: Any,

pool: Optional[concurrent.futures.ProcessPoolExecutor] = None,

) -> None:

"""Runs the full inference process on a batch of ZMWs and writes to fastq.

Args:

inputs: Iterable of inputs, one for each ZMW, each of which has three

elements (name of zmw, aligned_subreads, DcConfig).

model: An initialized model that will be used to make predictions.

model_params: Parameters for the model.

output_writer: File writer where fastq or bam output will be written.

options: Some options that apply to various stages of the inference run.

batch_name: Name of batch used for runtime metrics.

outcome_counter: Counts outcomes for each ZMW.

stats_counter: Global counter to gather example statistics.

pool: Process pool to run the preprocessing on. If None or empty,

preprocessing will be done sequentially on the main process.

"""

before_batch = time.time()

if options.cpus == 0:

# Preprocess ZMWs one at a time in the main process without multiprocessing.

outputs = [preprocess(one_zmw=one_zmw) for one_zmw in inputs] # pytype: disable=wrong-arg-types # always-use-return-annotations

else:

assert pool

# Each call to preprocess gets one ZMW from inputs.

outputs = list(pool.map(preprocess, inputs))

feature_dicts_for_zmws, counters = zip(*outputs)

num_zmws = len(feature_dicts_for_zmws)

for counter in counters:

stats_counter += counter

batch_total_examples = sum([len(zmw) for zmw in feature_dicts_for_zmws])

batch_total_subreads = sum([len(subreads) for _, subreads, _, _ in inputs])

timelog(

stage='preprocess',

item=batch_name,

before=before_batch,

num_examples=batch_total_examples,

num_subreads=batch_total_subreads,

num_zmws=num_zmws,

)

if FLAGS.end_after_stage in [DebugStage.TF_EXAMPLES, DebugStage.DC_INPUT]:

return

before = time.time()

before_skipping = time.time()

feature_dicts_for_model = []

predictions_for_skipped_windows = []

for one_zmw in feature_dicts_for_zmws:

for window in one_zmw:

skip_example = False

# Window can be skipped because it is larger than max example width.

if window['overflow']:

dc_output_for_window = process_skipped_window(window, options)

predictions_for_skipped_windows.append(dc_output_for_window)

skip_example = True

# Or window can be skipped if skip_windows_above is set and average base

# quality score is above the threshold.

if options.skip_windows_above and not skip_example:

avg_ccs_base_quality = utils.avg_phred(

window['ccs_base_quality_scores']

)

if avg_ccs_base_quality > options.skip_windows_above:

dc_output_for_window = process_skipped_window(window, options)

predictions_for_skipped_windows.append(dc_output_for_window)

skip_example = True

if not skip_example:

feature_dicts_for_model.append(window)

time_to_skip = time.time() - before_skipping

before_run_model = time.time()

predictions_from_model = run_model_on_examples(

feature_dicts_for_model, model, model_params, options

)

time_to_run_model = time.time() - before_run_model

predictions = predictions_from_model + predictions_for_skipped_windows

def percent_of_examples(numerator):

if not predictions:

return 0 # Avoid dividing by zero.

return 100 * (numerator / len(predictions))

logging.info(

(

'Example summary: ran model=%d (%0.2f%%; %0.3fs) skip=%d (%0.2f%%;'

' %0.3fs) total=%d.'

),

len(predictions_from_model),

percent_of_examples(len(predictions_from_model)),

time_to_run_model,

len(predictions_for_skipped_windows),

percent_of_examples(len(predictions_for_skipped_windows)),

time_to_skip,

len(predictions),

)

timelog(

stage='run_model',

item=batch_name,

before=before,

num_examples=batch_total_examples,

num_subreads=batch_total_subreads,

num_zmws=num_zmws,

)

if FLAGS.end_after_stage == DebugStage.RUN_MODEL:

return

before = time.time()

# Sort predictions prior to grouping

# pylint: disable=g-long-lambda

predictions = sorted(

predictions, key=lambda dc: (dc.molecule_name, dc.window_pos)

)

for zmw, predictions_for_zmw in itertools.groupby(

predictions, lambda p: p.molecule_name

):

predictions_for_zmw = list(predictions_for_zmw)

fastq_string = stitch_utils.stitch_to_fastq(

molecule_name=zmw,

predictions=predictions_for_zmw,

max_length=options.max_length,

min_quality=options.min_quality,

min_length=options.min_length,

outcome_counter=outcome_counter,

)

if fastq_string:

# FASTQs are written with gfile, bams are written with pysam.

if isinstance(output_writer, gfile.GFile):

output_writer.write(fastq_string)

else:

name, seq, _, qual = fastq_string.splitlines()

# Remove the @ prefix from sequence name.

name = name[1:]

record = pysam.AlignedSegment()

record.query_name = name

record.query_sequence = seq

record.query_qualities = pysam.qualitystring_to_array(qual)

record.flag = 4 # unmapped.

record.mapping_quality = 255

zmw = int(name.split('/')[1])

record.set_tags([

('ec', predictions_for_zmw[0].ec or -1, 'f'),

('np', predictions_for_zmw[0].np_num_passes, 'i'),

('rq', predictions_for_zmw[0].rq, 'f'),

('RG', predictions_for_zmw[0].rg, 'Z'),

('zm', zmw, 'i'),

])

output_writer.write(record)

timelog(

stage='stitch_and_write_fastq',

item=batch_name,

before=before,

num_examples=batch_total_examples,

num_subreads=batch_total_subreads,

num_zmws=num_zmws,

)

logging.info(

'Processed a batch of %d ZMWs in %0.3f seconds',

len(inputs),

time.time() - before_batch,

"""Save CSV of runtime."""

df = pd.DataFrame(time_points)

# Save CSV to file.

with tf.io.gfile.GFile(f'{output_prefix}.csv', 'w') as writer:

"""Output statistics into a file."""

json_stats = json.dumps(counter, indent=True)

with tf.io.gfile.GFile(f'{output_prefix}.json', 'w') as writer:

"""Performs an inference run."""

# Determine if --checkpoint is a saved model.

use_saved_model = tf.io.gfile.exists(FLAGS.checkpoint) and tf.io.gfile.exists(

f'{FLAGS.checkpoint}/saved_model.pb'

)

# Load model parameters

params = model_utils.read_params_from_json(checkpoint_path=FLAGS.checkpoint)

dc_config = pre_lib.DcConfig(

params.max_passes,

params.max_length,

params.use_ccs_bq,

)

if not FLAGS.max_base_quality:

raise ValueError(

'--max_base_quality should be set to a valid number. If unset bases'

' with error probability of 0, will result in invalid quality scores.'

)

# Attempt to read default calibration values from model params.json.

# If not found, set to 'skip'.

if FLAGS.dc_calibration is None:

dc_calibration_values = params.get('dc_calibration', 'skip')

if dc_calibration_values != 'skip':

logging.info(

(

'DeepConsensus base calibration values read from '

'model params.json: %s'

),

dc_calibration_values,

)

else:

dc_calibration_values = FLAGS.dc_calibration

dc_calibration_values = calibration_lib.parse_calibration_string(

dc_calibration_values

)

if not FLAGS.ccs_calibration:

raise ValueError(

'--ccs_calibration should be set to "skip" '

'or to base calibration scores.'

)

ccs_calibration_values = calibration_lib.parse_calibration_string(

FLAGS.ccs_calibration

)

options = InferenceOptions(

max_length=params.max_length,

example_height=dc_config.tensor_height,

max_passes=params.max_passes,

min_quality=FLAGS.min_quality,

min_length=FLAGS.min_length,

batch_size=FLAGS.batch_size,

cpus=FLAGS.cpus,

skip_windows_above=FLAGS.skip_windows_above,

use_saved_model=use_saved_model,

max_base_quality=FLAGS.max_base_quality,

dc_calibration_values=dc_calibration_values,

use_ccs_bq=params.use_ccs_bq,

ccs_calibration_values=ccs_calibration_values,

)

outcome_counter = stitch_utils.OutcomeCounter()

stats_counter = collections.Counter()

pool = None

if options.cpus > 0:

# Spin up multiple processes, each taking the next ZMW when ready.

pool = concurrent.futures.ProcessPoolExecutor(max_workers=options.cpus)

logging.info('Using multiprocessing: cpus is %s.', options.cpus)

elif options.cpus < 0:

raise ValueError(

'Number of processes must be positive '

'(for multiprocessing) or 0 (for serial execution).'

)

# Set up model.

before_model_setup = time.time()

loaded_model, model_params = initialize_model(

checkpoint_path=FLAGS.checkpoint, params=params, options=options

)

logging.info('Model setup took %s seconds.', time.time() - before_model_setup)

# Initialize output fastq writer.

output_fname = FLAGS.output

correct_suffix = (

output_fname.endswith('.fq')

or output_fname.endswith('.fastq')

or output_fname.endswith('.bam')

)

if not correct_suffix:

raise NameError('Filename must end in .fq, .fastq, or .bam')

output_dir = os.path.dirname(output_fname)

if not tf.io.gfile.exists(output_dir):

tf.io.gfile.makedirs(output_dir)

if output_fname.endswith('.fq') or output_fname.endswith('.fastq'):

output_writer = gfile.Open(output_fname, 'wb')

else:

ccs_bam_header = pysam.AlignmentFile(FLAGS.ccs_bam, check_sq=False).header

output_writer = pysam.AlignmentFile(

output_fname, 'wb', header=ccs_bam_header

)

input_file_generator = stream_bam(

subreads_to_ccs=FLAGS.subreads_to_ccs,

ccs_bam=FLAGS.ccs_bam,

options=options,

)

num_zmws_to_batch = FLAGS.batch_zmws

# Process ZMWs.

before_all_zmws = time.time()

zmw_counter = 0

batch_count = 0

stored_n_zmws = []

for zmw, subreads, dc_config, window_widths in input_file_generator:

if FLAGS.limit and zmw_counter >= FLAGS.limit:

break

zmw_counter += 1

stored_n_zmws.append((zmw, subreads, dc_config, window_widths))

if num_zmws_to_batch and len(stored_n_zmws) >= num_zmws_to_batch:

inference_on_n_zmws(

inputs=stored_n_zmws,

model=loaded_model,

model_params=model_params,

output_writer=output_writer,

options=options,

batch_name=str(batch_count),

outcome_counter=outcome_counter,

stats_counter=stats_counter,

pool=pool,

)

batch_count += 1

stored_n_zmws = []

logging.info(

'Processed %s ZMWs in %0.3f seconds',

zmw_counter,

time.time() - before_all_zmws,

)

if stored_n_zmws:

inference_on_n_zmws(

inputs=stored_n_zmws,

model=loaded_model,

model_params=model_params,

output_writer=output_writer,

options=options,

batch_name=str(batch_count),

outcome_counter=outcome_counter,

stats_counter=stats_counter,

pool=pool,

)

if pool:

pool.shutdown(wait=True)

output_writer.close()

logging.info(

'Processed %s ZMWs in %0.3f seconds',

zmw_counter,

time.time() - before_all_zmws,

)

logging.info('Outcome counts: %s', outcome_counter)

save_runtime(time_points=timing, output_prefix=f'{output_fname}.runtime')

save_counters(stats_counter, output_prefix=f'{output_fname}.inference')

"""Main entry point."""

if FLAGS.ccs_fasta:

raise NotImplementedError(

'The --ccs_fasta flag has been deprecated. '

'Please use --ccs_bam instead.'

)

if FLAGS.use_only_gpu_index:

with tf.device(f'GPU:{FLAGS.use_only_gpu_index}'):

outcome_counter = run()

else:

outcome_counter = run()

if not outcome_counter.success: