viral_complexity / utils.py
Hack90's picture
Update utils.py
61488b4 verified
from collections import namedtuple
import numpy as np
from scipy.interpolate import interp1d
import torch
import matplotlib.pyplot as plt
# Mapping of nucleotides to float coordinates
mapping_easy = {
"A": np.array([0.5, -0.8660254037844386]),
"T": np.array([0.5, 0.8660254037844386]),
"G": np.array([0.8660254037844386, -0.5]),
"C": np.array([0.8660254037844386, 0.5]),
"N": np.array([0, 0]),
}
# coordinates for x+iy
Coord = namedtuple("Coord", ["x", "y"])
# coordinates for a CGR encoding
CGRCoords = namedtuple("CGRCoords", ["N", "x", "y"])
# coordinates for each nucleotide in the 2d-plane
DEFAULT_COORDS = {"A": Coord(1, 1), "C": Coord(-1, 1), "G": Coord(-1, -1), "T": Coord(1, -1)}
# Function to convert a DNA sequence to a list of coordinates
def _dna_to_coordinates(dna_sequence: str, mapping: dict[str, np.ndarray]) -> np.ndarray:
dna_sequence = dna_sequence.upper()
coordinates = np.array([mapping.get(nucleotide, mapping["N"]) for nucleotide in dna_sequence])
return coordinates
# Function to create the cumulative sum of a list of coordinates
def _get_cumulative_coords(mapped_coords):
cumulative_coords = np.cumsum(mapped_coords, axis=0)
return cumulative_coords
def generate_2d_sequence(seq):
dna_sequence = seq.upper()
mapped_coords = _dna_to_coordinates(dna_sequence, mapping_easy)
cumulative_coords = _get_cumulative_coords(mapped_coords)
# Scale the input data using standardization
x_train = cumulative_coords[:, 0]
y_train = cumulative_coords[:, 1]
x_train_scaled = (x_train - x_train.mean()) / x_train.std()
y_train_scaled = (y_train - y_train.mean()) / y_train.std()
scaled_coords = np.column_stack((x_train_scaled, y_train_scaled))
# example["2D_Sequence"] = cumulative_coords.tolist()
# example["2D_Sequence_Scaled"] = scaled_coords.tolist()
# Interpolate the 2D sequences to have exactly 1000 pairs
interpolated_coords = y_train_scaled # default to filter out bad examples
if len(scaled_coords) != 1000:
try:
t = np.linspace(0, 1, len(scaled_coords))
t_new = np.linspace(0, 1, 1000)
interp_func_x = interp1d(t, scaled_coords[:, 0], kind="linear")
interp_func_y = interp1d(t, scaled_coords[:, 1], kind="linear")
interpolated_coords = interp_func_x(t_new)
except Exception as e:
print(f"Interpolation error: {e}")
tensor_2d_rep_y = torch.Tensor(interpolated_coords).reshape(1,1000)
return y_train_scaled, x_train_scaled
def generate_2d_sequence_small(seq):
dna_sequence = seq.upper()
mapped_coords = _dna_to_coordinates(dna_sequence, mapping_easy)
cumulative_coords = _get_cumulative_coords(mapped_coords)
# Scale the input data using standardization
x_train = cumulative_coords[:, 0]
y_train = cumulative_coords[:, 1]
x_train_scaled = (x_train - x_train.mean()) / x_train.std()
y_train_scaled = (y_train - y_train.mean()) / y_train.std()
scaled_coords = np.column_stack((x_train_scaled, y_train_scaled))
# example["2D_Sequence"] = cumulative_coords.tolist()
# example["2D_Sequence_Scaled"] = scaled_coords.tolist()
# Interpolate the 2D sequences to have exactly 1000 pairs
interpolated_coords = y_train_scaled # default to filter out bad examples
if len(scaled_coords) != 1000:
try:
t = np.linspace(0, 1, len(scaled_coords))
t_new = np.linspace(0, 1, 400)
interp_func_x = interp1d(t, scaled_coords[:, 0], kind="linear")
interp_func_y = interp1d(t, scaled_coords[:, 1], kind="linear")
interpolated_coords = interp_func_y(t_new)
except Exception as e:
print(f"Interpolation error: {e}")
tensor_2d_rep_y = torch.Tensor(interpolated_coords).reshape(400)
return tensor_2d_rep_y
def plot_seq_full_label(df, filter):
ncols = len(filter)
unique_ids = df.label_id.unique()
print(unique_ids)
unique_ids_plot = [id for id in unique_ids if id in filter]
print(unique_ids_plot)
fig, axs = plt.subplots(ncols=ncols)
for i, id in enumerate(unique_ids_plot):
# data = (df[df['label_id'] == id].sample((n=3)))['seq'].values[0]
# print(data)
data = generate_2d_sequence_small(df[df['label_id'] == id].sample(n=1)['seq'].values[0]).numpy()
# two_d = generate_2d_sequence(data)[0]
axs[i].plot(data)
return fig