<jupyter_start><jupyter_text>Preprocessing<jupyter_code>split = 'test'
# 1.25 M pos/neg/unk pairs per charge, similar to pairs_dist.
num_pairs = 3_750_000 * 4
# Metadata table for all (identified and unidentified) spectra.   
spectrum_info = pd.merge(
    pd.read_parquet(os.path.join(
        os.environ['GLEAMS_HOME'], 'data', 'embed',
        f'embed_{config.massivekb_task_id}_{split}.parquet')),
    pd.read_parquet(os.path.join(
        os.environ['GLEAMS_HOME'], 'data', 'feature',
        f'feature_{config.massivekb_task_id}_{split}.parquet'))
    [['dataset', 'filename', 'scan', 'sequence']],
    'left', ['dataset', 'filename', 'scan'], copy=False)
spectrum_info.to_parquet('dot_embed_metadata.parquet', index=False)
# Generate all types of pairs (positive, negative, unknown).
metadata.generate_pairs_positive('dot_embed_metadata.parquet', config.charges)
metadata.generate_pairs_negative('dot_embed_metadata.parquet', config.charges,
                                 config.pair_mz_tolerance,
                                 config.negative_pair_fragment_tolerance,
                                 config.negative_pair_matching_fragments_threshold)
max_spectra = 10_000_000
generate_pairs_unknown('dot_embed_metadata.parquet', config.charges,
                       config.pair_mz_tolerance, max_spectra)
pair_generator = PairIndexSequence(
    [f'dot_embed_metadata_pairs_pos_{charge}.npy'
     for charge in np.arange(config.charges[0], config.charges[1] + 1)],
    [f'dot_embed_metadata_pairs_neg_{charge}.npy'
     for charge in np.arange(config.charges[0], config.charges[1] + 1)],
    [f'dot_embed_metadata_pairs_unk_{charge}.npy'
     for charge in np.arange(config.charges[0], config.charges[1] + 1)],
    config.batch_size, num_pairs, False)
# Subset the spectrum info to only include the selected spectra from the pair
# generator.
spectrum_indexes = np.hstack((np.vstack(pair_generator.pairs_pos).reshape(-1),
                              np.vstack(pair_generator.pairs_neg).reshape(-1),
                              np.vstack(pair_generator.pairs_unk).reshape(-1)))
spectrum_info = spectrum_info.loc[np.unique(spectrum_indexes)]<jupyter_output><empty_output><jupyter_text>Dot product<jupyter_code># Read the selected spectra from the peak files.
dataset_total = spectrum_info['dataset'].nunique()
spectra = []
for dataset_i, (dataset, md_dataset) in enumerate(
        spectrum_info.groupby('dataset', sort=False), 1):
    logging.info('Process dataset %s (%d files) [%3d/%3d]', dataset,
                 md_dataset['filename'].nunique(), dataset_i, dataset_total)
    spectra.extend(joblib.Parallel(n_jobs=-1, backend='multiprocessing')(
        joblib.delayed(_get_spectra_from_file)(dataset, filename,
                                               md_file['scan'])
        for filename, md_file in md_dataset.groupby(
            'filename', sort=False)))
spectra = collections.ChainMap(*spectra)
# Compute the dot products (high/low resolution) between all spectra pairs.
spectra_arr = [], []
for batch_i in tqdm.tqdm(range(len(pair_generator)),
                         desc='Spectra converted', unit='batch'):
    (batch_i1, batch_i2), _ = pair_generator[batch_i]
    for pair_i, arr_i in zip(itertools.chain(batch_i1, batch_i2),
                             np.hstack((np.zeros_like(batch_i1),
                                        np.ones_like(batch_i2)))):
        spec = spectra[f"{spectrum_info.at[pair_i, 'dataset']}/"
                       f"{spectrum_info.at[pair_i, 'filename']}/"
                       f"{spectrum_info.at[pair_i, 'scan']}"]
        spectra_arr[arr_i].append(
            np.pad([spec.mz, spec.intensity],
                   ((0, 0), (config.max_peaks_used - len(spec.mz), 0)),
                   'constant'))

fragment_mz_tol_high_res, fragment_mz_tol_low_res = 0.05, 0.8
spectra_arr1 = np.asarray(spectra_arr[0])
spectra_arr2 = np.asarray(spectra_arr[1])
logger.info('Compute dot product (high resolution; '
            'fragment m/z tolerance = %.2f)', fragment_mz_tol_high_res)
dot_high_res = dot(
    spectra_arr1, spectra_arr2, np.zeros(spectra_arr1.shape[0], np.float32),
    fragment_mz_tol_high_res)
logger.info('Compute dot product (low resolution; '
            'fragment m/z tolerance = %.2f)', fragment_mz_tol_low_res)
dot_low_res = dot(
    spectra_arr1, spectra_arr2, np.zeros(spectra_arr1.shape[0], np.float32),
    fragment_mz_tol_low_res)<jupyter_output><empty_output><jupyter_text>GLEAMS Euclidean distance<jupyter_code># Compute the Euclidean distances between all embeddings pairs.
embeddings = np.load(os.path.join(
    os.environ['GLEAMS_HOME'], 'data', 'embed',
    f'embed_{config.massivekb_task_id}_{split}.npy'))

scores_embed = []
for batch_i in tqdm.tqdm(range(len(pair_generator)),
                         desc='GLEAMS distances calculated', unit='batch'):
    (batch_i1, batch_i2), _ = pair_generator[batch_i]
    for pair1, pair2 in zip(batch_i1, batch_i2):
        scores_embed.append(ssd.euclidean(embeddings[pair1],
                                          embeddings[pair2]))<jupyter_output><empty_output><jupyter_text>Dot product versus GLEAMS Euclidean distance<jupyter_code>labels = np.hstack([pair_generator[batch_i][1]
                    for batch_i in range(len(pair_generator))])
dot_embed = pd.DataFrame({'dot_low_res': dot_low_res,
                          'dot_high_res': dot_high_res,
                          'gleams_dist': scores_embed,
                          'pair_type': labels}).sort_values('pair_type')
# Convert pair type to nice labels.
dot_embed['pair_type'] = dot_embed['pair_type'].map(
    {1: 'Positive', -1: 'Negative', 0: 'Unknown'})
# Add precursor information.
spec_idx = np.hstack([pair_generator[batch_i][0][0]
                      for batch_i in range(len(pair_generator))])
dot_embed[['charge', 'mz']] = (spectrum_info.loc[spec_idx, ['charge', 'mz']]
                               .reset_index(drop=True))
dot_embed.to_parquet('dot_embed.parquet', index=False)
# dot_embed = pd.read_parquet('dot_embed.parquet')
width = 7

jg = sns.jointplot(
    data=dot_embed, x='dot_high_res', y='gleams_dist', hue='pair_type',
    hue_order=['Positive', 'Negative', 'Unknown'], height=width, s=1,
    marker='.', rasterized=True, joint_kws={'alpha': 0.1})

jg.ax_joint.legend(title='Pair type')
jg.set_axis_labels('Spectrum dot product', 'GLEAMS euclidean distance')

plt.savefig('dot_embed_type.png', dpi=300, bbox_inches='tight')
plt.show()
plt.close()
height = 7

hue_order = (2, 3, 4, 5)
jg = sns.JointGrid(height=height)
sns.scatterplot(data=dot_embed, x='dot_high_res', y='gleams_dist',
                hue='charge', hue_order=hue_order, alpha=0.1, s=1, marker='.',
                rasterized=True, ax=jg.ax_joint)
sns.kdeplot(data=dot_embed, x='dot_high_res', hue='charge',
            hue_order=hue_order, legend=False, common_norm=False, fill=True,
            ax=jg.ax_marg_x)
sns.kdeplot(data=dot_embed, y='gleams_dist', hue='charge',
            hue_order=hue_order, legend=False, common_norm=False, fill=True,
            ax=jg.ax_marg_y)

jg.ax_joint.legend(loc='upper right', title='Precursor charge')
jg.set_axis_labels('Spectrum dot product', 'GLEAMS euclidean distance')

plt.savefig('dot_embed_charge.png', dpi=300, bbox_inches='tight')
plt.show()
plt.close()
height = 7

jg = sns.JointGrid(height=height)
sns.scatterplot(data=dot_embed, x='dot_high_res', y='gleams_dist',
                alpha=0.1, s=1, c=dot_embed['mz'], marker='.',
                cmap=plt.cm.get_cmap('YlGnBu'), rasterized=True,
                ax=jg.ax_joint)
sns.kdeplot(data=dot_embed, x='dot_high_res', color='black', legend=False,
            common_norm=False, fill=True, ax=jg.ax_marg_x)
sns.kdeplot(data=dot_embed, y='gleams_dist', color='black', legend=False,
            common_norm=False, fill=True, ax=jg.ax_marg_y)

ax_joint_pos = jg.ax_joint.get_position()
cbar_ax = jg.fig.add_axes([1.025, ax_joint_pos.x0 + 0.05,
                           0.025, ax_joint_pos.height - 0.1])
colorbar = jg.fig.colorbar(jg.ax_joint.get_children()[0], cax=cbar_ax)
colorbar.solids.set(alpha=1)
colorbar.set_label('Precursor m/z', size='large', labelpad=15)

jg.set_axis_labels('Spectrum dot product', 'GLEAMS euclidean distance')

plt.savefig('dot_embed_mz.png', dpi=300, bbox_inches='tight')
plt.show()
plt.close()
logging.shutdown()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1383422778771616, 'fraction_numerical': 0.014788312462731067, 'mean_word_length': 3.653718091009989, 'pattern_counts': {'":': 0, '<': 17, '<?xml version=': 0, '>': 17, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 8, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '4', 'dir': 'jupyter-structured-clean-dedup', 'id': '467465', 'license': 'BSD-3-Clause', 'n_tokens_mistral': 3032, 'n_tokens_neox': 2886, 'n_words': 469, 'path': 'dot_embed.ipynb', 'repo_name': 'bittremieux/GLEAMS_notebooks', 'size': '8385'}