def test_write_spectra_to_sdf(self):
"""Tests predicting and writing spectra to file."""
predictor = spectra_predictor.NeimsSpectraPredictor(model_checkpoint_dir="")
mols_from_file = spectra_predictor.get_mol_list_from_sdf(
self.test_file_short)
fingerprints, mol_weights = predictor.get_inputs_for_model_from_mol_list(
mols_from_file)
spectra_predictions = predictor.make_spectra_prediction(
fingerprints, mol_weights)
spectra_predictor.update_mols_with_spectra(mols_from_file,
spectra_predictions)
_, fpath = tempfile.mkstemp(dir=self.temp_dir)
spectra_predictor.write_rdkit_mols_to_sdf(mols_from_file, fpath)
# Test contents of newly written file:
new_mol_list = spectra_predictor.get_mol_list_from_sdf(fpath)
for idx, mol in enumerate(new_mol_list):
peak_string_from_file = mol.GetProp(
spectra_predictor.PREDICTED_SPECTRA_PROP_NAME)
peak_locs, peak_intensities = feature_utils.parse_peaks(
peak_string_from_file)
dense_mass_spectra = feature_utils.make_dense_mass_spectra(
peak_locs, peak_intensities, ms_constants.MAX_PEAK_LOC)
self.assertSequenceAlmostEqual(
dense_mass_spectra, spectra_predictions[idx, :].astype(np.int32),
delta=1.)
def make_spectra_array(mol_list):
"""Grab spectra pertaining to same molecule in one np.array.
Args:
mol_list: list of rdkit.Mol objects. Each Mol should contain
information about the spectra, as stored in NIST.
Output:
np.array of spectra of shape (number of spectra, max spectra length)
"""
mass_spec_spectra = np.zeros( ( len(mol_list), ms_constants.MAX_PEAK_LOC))
for idx, mol in enumerate(mol_list):
spectra_str = mol.GetProp(ms_constants.SDF_TAG_MASS_SPEC_PEAKS)
spectral_locs, spectral_intensities = feature_utils.parse_peaks(spectra_str)
dense_mass_spec = feature_utils.make_dense_mass_spectra(
spectral_locs, spectral_intensities, ms_constants.MAX_PEAK_LOC)
mass_spec_spectra[idx, :] = dense_mass_spec
return mass_spec_spectra
def make_mol_dict(mol,
max_atoms=ms_constants.MAX_ATOMS,
max_mass_spec_peak_loc=ms_constants.MAX_PEAK_LOC,
add_hs_to_molecule=ms_constants.ADD_HS_TO_MOLECULES):
"""Place molecule information from rdkit.Mol.Prop in a dictionary.
Reads information from an rdkit mol with the desired tags for reading into
a dictionary.
Args:
mol : An rdkit.Mol object read from an sdf that contains sdf tags:
fmap_constants.NAME, 'INCHIKEY', 'MASS SPECTRAL PEAKS'
to be converted into dictionary keys. These tags should be stored as
properties by the rdkit.Mol
max_atoms : maximum number of allowed atoms.
max_mass_spec_peak_loc : largest mass/charge ratio to allow in a spectra
add_hs_to_molecule : Whehter or not to include hydrogen atoms in the
molecule
Returns:
mol_dict : A dictionary with following molecule info
keys : name, inchi_key, circular_fp_1024, molecule_weight, smiles,
atom_weights, mass_spec_dense
Notes:
- smiles recorded in mol_dict here have been canonicalized, and may not
correspond with the order in the original SDF
- atom_weights vector are length max_atoms
- circular_fp_1024 has 1024 bits, with default radius of 2.
"""
# For canonicalizing molecules
smiles_canon = feature_utils.get_smiles_string(mol)
mol_formula = feature_utils.get_molecular_formula(mol)
mol_canon = Chem.MolFromSmiles(smiles_canon)
if add_hs_to_molecule:
mol_canon = Chem.AddHs(mol_canon)
mass_spec_locs, mass_spec_intensities = feature_utils.parse_peaks(
mol.GetProp(ms_constants.SDF_TAG_MASS_SPEC_PEAKS))
mass_spec_dense_vec = feature_utils.make_dense_mass_spectra(
mass_spec_locs, mass_spec_intensities, max_mass_spec_peak_loc)
atom_wts = feature_utils.get_padded_atom_weights(mol_canon, max_atoms)
atom_ids = feature_utils.get_padded_atom_ids(mol_canon, max_atoms)
adjacency_matrix = feature_utils.get_padded_adjacency_matrix(
mol_canon, max_atoms)
mol_dict = {
fmap_constants.NAME:
mol.GetProp(ms_constants.SDF_TAG_NAME),
fmap_constants.INCHIKEY:
mol.GetProp(ms_constants.SDF_TAG_INCHIKEY),
fmap_constants.MOLECULAR_FORMULA:
mol_formula,
fmap_constants.MOLECULE_WEIGHT:
float(mol.GetProp(ms_constants.SDF_TAG_MOLECULE_MASS)),
fmap_constants.SMILES:
smiles_canon,
fmap_constants.ATOM_WEIGHTS:
atom_wts,
fmap_constants.ATOM_IDS:
atom_ids,
fmap_constants.DENSE_MASS_SPEC:
mass_spec_dense_vec,
fmap_constants.ADJACENCY_MATRIX:
adjacency_matrix,
}
# Making fingerprint features:
mol_dict.update(feature_utils.all_circular_fingerprints_to_dict(mol))
return mol_dict
def setUp(self):
super(ParseSdfUtilsTest, self).setUp()
self.test_data_directory = test_utils.test_dir('testdata/')
self.test_file_long = os.path.join(self.test_data_directory,
'test_14_mend.sdf')
self.test_file_short = os.path.join(self.test_data_directory,
'test_2_mend.sdf')
self.temp_dir = tempfile.mkdtemp(dir=absltest.get_default_test_tmpdir())
# Expected result for list of molecule dicts
self.expected_mol_dicts = [{
fmap_constants.NAME: 'Methane, diazo-',
fmap_constants.INCHIKEY: 'YXHKONLOYHBTNS-UHFFFAOYSA-N',
fmap_constants.MOLECULAR_FORMULA: 'CH2N2',
fmap_constants.SMILES: 'C=[N+]=[N-]',
'parsed_smiles': [28, 18, 81, 51, 4, 83, 18, 81, 51, 5, 83],
fmap_constants.SMILES_TOKEN_LIST_LENGTH: 11
}, {
fmap_constants.NAME: (
'(4-(4-Chlorphenyl)-3-morpholino-pyrrol-2-yl)-butenedioic acid,'
' dimethyl ester'),
fmap_constants.INCHIKEY:
'PNYUDNYAXSEACV-RVDMUPIBSA-N',
fmap_constants.MOLECULAR_FORMULA:
'C20H21ClN2O5',
fmap_constants.SMILES:
'COC(=O)/C=C(/C(=O)OC)c1[nH]cc(-c2ccc(Cl)cc2)c1N1CCOCC1',
'parsed_smiles': [
28, 55, 28, 2, 18, 55, 3, 7, 28, 18, 28, 2, 7, 28, 2, 18, 55, 3, 55,
28, 3, 84, 9, 81, 85, 40, 83, 84, 84, 2, 5, 84, 10, 84, 84, 84, 2,
31, 3, 84, 84, 10, 3, 84, 9, 51, 9, 28, 28, 55, 28, 28, 9
],
fmap_constants.SMILES_TOKEN_LIST_LENGTH:
53,
}]
for mol_dict in self.expected_mol_dicts:
token_arr = mol_dict['parsed_smiles']
sequence_length = mol_dict[
fmap_constants.SMILES_TOKEN_LIST_LENGTH]
mol_dict['parsed_smiles'] = np.pad(
token_arr, (0, ms_constants.MAX_TOKEN_LIST_LENGTH - sequence_length),
'constant')
mol_weights = [42.0217981, 404.1139]
atom_weights_list = [[12.011, 14.007, 14.007], [
12.011, 15.999, 12.011, 15.999, 12.011, 12.011, 12.011, 15.999, 15.999,
12.011, 12.011, 14.007, 12.011, 12.011, 12.011, 12.011, 12.011, 12.011,
35.453, 12.011, 12.011, 12.011, 14.007, 12.011, 12.011, 15.999, 12.011,
12.011
]]
atom_ids_list = [[6, 7, 7], [
6, 8, 6, 8, 6, 6, 6, 8, 8, 6, 6, 7, 6, 6, 6, 6, 6, 6, 17, 6, 6, 6, 7, 6,
6, 8, 6, 6
]]
adjacency_matrix_list = [
np.array(
[
0., 2., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 2., 0.,
2., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 2., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0.
],
dtype='int32'),
np.array(
[
0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0.,
1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 2.,
1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 2., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 2., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 2., 0., 1., 0., 0., 0.,
1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 2., 1., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 2., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 1., 0., 0., 1., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 1., 0., 0., 0., 0., 0., 4., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 4., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 4., 0., 4., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 4., 0., 4., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 4., 0., 1., 0., 0., 0., 0., 0., 0., 4., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 1., 0., 4., 0., 0., 0., 0., 4., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
4., 0., 4., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 4.,
0., 4., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 4., 0.,
1., 4., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 4., 0., 0., 4., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 4., 0., 0., 0., 0., 4., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
4., 0., 0., 4., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 1., 0., 0., 0., 1.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 1., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 1., 0., 1., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 1., 0., 1., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 1., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
1., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0.
],
dtype='int32')
]
mass_spec_peak_locs = [[22, 23, 24, 25, 26, 27, 28, 30, 31, 32], [
32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 49, 50,
51
]]
mass_spec_peak_intensities = [[
110, 220, 999, 25, 12, 58, 179, 22, 110, 425
], [
12, 7, 28, 999, 57, 302, 975, 8, 53, 176, 99, 122, 117, 155, 9, 7, 6,
28, 59
]]
# Special hparams so that expected arrays can be smaller.
self.hparams = tf.contrib.training.HParams(
intensity_power=1.0,
max_atoms=30,
max_mass_spec_peak_loc=60,
eval_batch_size=len(self.expected_mol_dicts))
for i in range(len(self.expected_mol_dicts)):
self.expected_mol_dicts[i][
fmap_constants.MOLECULE_WEIGHT] = mol_weights[i]
self.expected_mol_dicts[i][
fmap_constants.ATOM_WEIGHTS] = np.pad(
np.array(atom_weights_list[i]),
(0, self.hparams.max_atoms - len(atom_weights_list[i])),
'constant')
self.expected_mol_dicts[i][fmap_constants.ATOM_IDS] = np.pad(
np.array(atom_ids_list[i]),
(0, self.hparams.max_atoms - len(atom_ids_list[i])), 'constant')
self.expected_mol_dicts[i][fmap_constants.ADJACENCY_MATRIX] = (
adjacency_matrix_list[i])
self.expected_mol_dicts[i][fmap_constants.DENSE_MASS_SPEC] = (
feature_utils.make_dense_mass_spectra(
mass_spec_peak_locs[i], mass_spec_peak_intensities[i],
self.hparams.max_mass_spec_peak_loc))