def bpnet_contrib(
model_dir,
output_file,
method="grad",
dataspec=None,
regions=None,
fasta_file=None, # alternative to dataspec
shuffle_seq=False,
shuffle_regions=False,
max_regions=None,
# reference='zeroes', # Currently the only option
# peak_width=1000, # automatically inferred from 'config.gin.json'
# seq_width=None,
contrib_wildcard='*/profile/wn,*/counts/pre-act', # specifies which contrib. scores to compute
batch_size=512,
gpu=0,
memfrac_gpu=0.45,
num_workers=10,
storage_chunk_size=512,
exclude_chr='',
include_chr='',
overwrite=False,
skip_bias=False):
"""Run contribution scores for a BPNet model
"""
from bpnet.extractors import _chrom_sizes
add_file_logging(os.path.dirname(output_file), logger, 'bpnet-contrib')
if gpu is not None:
create_tf_session(gpu, per_process_gpu_memory_fraction=memfrac_gpu)
else:
# Don't use any GPU's
os.environ['CUDA_VISIBLE_DEVICES'] = ''
if os.path.exists(output_file):
if overwrite:
os.remove(output_file)
else:
raise ValueError(
f"File exists {output_file}. Use overwrite=True to overwrite it"
)
config = read_json(os.path.join(model_dir, 'config.gin.json'))
seq_width = config['seq_width']
peak_width = config['seq_width']
# NOTE - seq_width has to be the same for the input and the target
#
# infer from the command line
# if seq_width is None:
# logger.info("Using seq_width = peak_width")
# seq_width = peak_width
# # make sure these are int's
# seq_width = int(seq_width)
# peak_width = int(peak_width)
# Split
contrib_wildcards = contrib_wildcard.split(",")
# Allow chr inclusion / exclusion
if exclude_chr:
exclude_chr = exclude_chr.split(",")
else:
exclude_chr = None
if include_chr:
include_chr = include_chr.split(",")
else:
include_chr = None
logger.info("Loading the config files")
model_dir = Path(model_dir)
logger.info("Creating the dataset")
from bpnet.datasets import StrandedProfile, SeqClassification
if fasta_file is not None:
if regions is None:
raise ValueError(
"fasta_file specified. Expecting regions to be specified as well"
)
dl_valid = SeqClassification(
fasta_file=fasta_file,
intervals_file=regions,
incl_chromosomes=include_chr,
excl_chromosomes=exclude_chr,
auto_resize_len=seq_width,
)
chrom_sizes = _chrom_sizes(fasta_file)
else:
if dataspec is None:
logger.info("Using dataspec used to train the model")
# Specify dataspec
dataspec = model_dir / "dataspec.yml"
ds = DataSpec.load(dataspec)
dl_valid = StrandedProfile(ds,
incl_chromosomes=include_chr,
excl_chromosomes=exclude_chr,
intervals_file=regions,
peak_width=peak_width,
shuffle=False,
seq_width=seq_width)
chrom_sizes = _chrom_sizes(ds.fasta_file)
# Setup contribution score trimming (not required currently)
if seq_width > peak_width:
# Trim
# make sure we can nicely trim the peak
logger.info("Trimming the output")
assert (seq_width - peak_width) % 2 == 0
trim_start = (seq_width - peak_width) // 2
trim_end = seq_width - trim_start
assert trim_end - trim_start == peak_width
elif seq_width == peak_width:
trim_start = 0
trim_end = peak_width
else:
raise ValueError("seq_width < peak_width")
seqmodel = SeqModel.from_mdir(model_dir)
# get all possible interpretation names
# make sure they match the specified glob
intp_names = [
name for name, _ in seqmodel.get_intp_tensors(preact_only=False)
if fnmatch_any(name, contrib_wildcards)
]
logger.info(f"Using the following interpretation targets:")
for n in intp_names:
print(n)
if max_regions is not None:
if len(dl_valid) > max_regions:
logging.info(
f"Using {max_regions} regions instead of the original {len(dl_valid)}"
)
else:
logging.info(
f"--max-regions={max_regions} is larger than the dataset size: {len(dl_valid)}. "
"Using the dataset size for max-regions")
max_regions = len(dl_valid)
else:
max_regions = len(dl_valid)
max_batches = np.ceil(max_regions / batch_size)
writer = HDF5BatchWriter(output_file, chunk_size=storage_chunk_size)
for i, batch in enumerate(
tqdm(dl_valid.batch_iter(batch_size=batch_size,
shuffle=shuffle_regions,
num_workers=num_workers),
total=max_batches)):
# store the original batch containing 'inputs' and 'targets'
if skip_bias:
batch['inputs'] = {
'seq': batch['inputs']['seq']
} # ignore all other inputs
if max_batches > 0:
if i > max_batches:
break
if shuffle_seq:
# Di-nucleotide shuffle the sequences
batch['inputs']['seq'] = onehot_dinucl_shuffle(
batch['inputs']['seq'])
for name in intp_names:
hyp_contrib = seqmodel.contrib_score(
batch['inputs']['seq'],
name=name,
method=method,
batch_size=None) # don't second-batch
# put contribution scores to the dictionary
# also trim the contribution scores appropriately so that
# the output will always be w.r.t. the peak center
batch[f"/hyp_contrib/{name}"] = hyp_contrib[:, trim_start:trim_end]
# trim the sequence as well
# Trim the sequence
batch['inputs']['seq'] = batch['inputs']['seq'][:, trim_start:trim_end]
# ? maybe it would it be better to have an explicit ContribFileWriter.
# that way the written schema would be fixed
writer.batch_write(batch)
# add chromosome sizes
writer.f.attrs['chrom_sizes'] = json.dumps(chrom_sizes)
writer.close()
logger.info(f"Done. Contribution score file was saved to: {output_file}")
def __init__(self, ds,
peak_width=200,
seq_width=None,
incl_chromosomes=None,
excl_chromosomes=None,
intervals_file=None,
intervals_format='bed',
include_metadata=True,
tasks=None,
include_classes=False,
shuffle=True,
interval_transformer=None,
track_transform=None,
total_count_transform=lambda x: np.log(1 + x)):
"""Dataset for loading the bigwigs and fastas
Args:
ds (bpnet.dataspecs.DataSpec): data specification containing the
fasta file, bed files and bigWig file paths
chromosomes (list of str): a list of chor
peak_width: resize the bed file to a certain width
intervals_file: if specified, use these regions to train the model.
If not specified, the regions are inferred from the dataspec.
intervals_format: interval_file format. Available: bed, bed3, bed3+labels
shuffle: True
track_transform: function to be applied to transform the tracks (shape=(batch, seqlen, channels))
total_count_transform: transform to apply to the total counts
TODO - shall we standardize this to have also the inverse operation?
"""
if isinstance(ds, str):
self.ds = DataSpec.load(ds)
else:
self.ds = ds
self.peak_width = peak_width
if seq_width is None:
self.seq_width = peak_width
else:
self.seq_width = seq_width
assert intervals_format in ['bed3', 'bed3+labels', 'bed']
self.shuffle = shuffle
self.intervals_file = intervals_file
self.intervals_format = intervals_format
self.incl_chromosomes = incl_chromosomes
self.excl_chromosomes = excl_chromosomes
self.total_count_transform = total_count_transform
self.track_transform = track_transform
self.include_classes = include_classes
# not specified yet
self.fasta_extractor = None
self.bw_extractors = None
self.bias_bw_extractors = None
self.include_metadata = include_metadata
self.interval_transformer = interval_transformer
# Load chromosome lengths
self.chrom_lens = _chrom_sizes(self.ds.fasta_file)
if self.intervals_file is None:
# concatenate the bed files
self.dfm = pd.concat([TsvReader(task_spec.peaks,
num_chr=False,
incl_chromosomes=incl_chromosomes,
excl_chromosomes=excl_chromosomes,
chromosome_lens=self.chrom_lens,
resize_width=max(self.peak_width, self.seq_width)
).df.iloc[:, :3].assign(task=task)
for task, task_spec in self.ds.task_specs.items()
if task_spec.peaks is not None])
assert list(self.dfm.columns)[:4] == [0, 1, 2, "task"]
if self.shuffle:
self.dfm = self.dfm.sample(frac=1)
self.tsv = None
self.dfm_tasks = None
else:
self.tsv = TsvReader(self.intervals_file,
num_chr=False,
# optional
label_dtype=int if self.intervals_format == 'bed3+labels' else None,
mask_ambigous=-1 if self.intervals_format == 'bed3+labels' else None,
# --------------------------------------------
incl_chromosomes=incl_chromosomes,
excl_chromosomes=excl_chromosomes,
chromosome_lens=self.chrom_lens,
resize_width=max(self.peak_width, self.seq_width)
)
if self.shuffle:
self.tsv.shuffle_inplace()
self.dfm = self.tsv.df # use the data-frame from tsv
self.dfm_tasks = self.tsv.get_target_names()
# remember the tasks
if tasks is None:
self.tasks = list(self.ds.task_specs)
else:
self.tasks = tasks
if self.include_classes:
assert self.dfm_tasks is not None
if self.dfm_tasks is not None:
assert set(self.tasks).issubset(self.dfm_tasks)
# setup bias maps per task
self.task_bias_tracks = {task: [bias for bias, spec in self.ds.bias_specs.items()
if task in spec.tasks]
for task in self.tasks}