Mols matrix to grid image

rdkit/Chem/Draw/UnitTestDraw.py

@@ -7,18 +7,15 @@
 #  which is included in the file license.txt, found at the root
 #  of the RDKit source tree.
 #
-import re
 import os
 import pathlib
+import re
 import sys
 import tempfile
 import unittest
 
 from rdkit import Chem
-from rdkit.Chem import AllChem
-from rdkit.Chem import Draw
-from rdkit.Chem import rdDepictor
-from rdkit.Chem import rdMolDescriptors
+from rdkit.Chem import AllChem, Draw, rdDepictor, rdMolDescriptors
 
 try:
   from rdkit.Chem.Draw import IPythonConsole
@@ -50,6 +47,51 @@ def setUp(self):
       IPythonConsole.UninstallIPythonRenderer()
     self.mol = Chem.MolFromSmiles('c1c(C[15NH3+])ccnc1[C@](Cl)(Br)[C@](Cl)(Br)F')
 
+    # For testMolsMatrixToLinear and testMolsMatrixToLinear (which test MolsMatrixToGridImage and its helper _MolsNestedToLinear)
+    s = "NC(C)C(=O)"
+    mol = Chem.MolFromSmiles(s)
+    natoms = mol.GetNumAtoms()
+    nbonds = mol.GetNumBonds()
+
+    # Set up matrix with oligomer count for the molecules
+    # Should produce this grid:
+    # NC(C)C(=O)
+    #                                NC(C)C(=O)NC(C)C(=O)
+    # NC(C)C(=O)NC(C)C(=O)NC(C)C(=O)                      NC(C)C(=O)NC(C)C(=O)NC(C)C(=O)NC(C)C(=O)
+    repeats = [[1], [0, 2], [3, 0, 4]]
+
+    # Create molecule if there are 1 or more oligomers;
+    # otherwise, use None for molecule because drawing is distorted if use Chem.MolFromSmiles("")
+    self.mols_matrix = [[Chem.MolFromSmiles(s * count) if count else None for count in row] for row in repeats]
+
+    self.legends_matrix = [[str(count) + " unit(s)" for count in row] for row in repeats]
+
+    def ith_item_list(nunits, items_per_unit, i=0):
+        return [((n * items_per_unit) + i) for n in range(nunits)]
+
+    self.highlightAtomLists_matrix = [[ith_item_list(count, natoms, 0) for count in row] for row in repeats]
+
+    # Another bond is created when molecule is oligomerized, so to keep the bond type consistent,
+    #   make items per unit one more than the number of bonds
+    self.highlightBondLists_matrix = [[ith_item_list(count, nbonds + 1, 1) for count in row] for row in repeats]
+
+    # Parametrize tests: In addition to supplying mols_matrix, supply 0-3 other matrices
+    # col labels: legends_matrix,       highlightAtomLists_matrix,      highlightBondLists_matrix      
+    #             Zero other matrices: 1 parameter set
+    self.param_sets = [(None,                None,                           None                          ), 
+                  # One other matrix: 3 parameter sets
+                  (self.legends_matrix, None,                           None                          ),
+                  (None,                self.highlightAtomLists_matrix, None                          ),
+                  (None,                None,                           self.highlightBondLists_matrix),
+                  # Two other matrices: 3 parameter sets
+                  (self.legends_matrix, self.highlightAtomLists_matrix, None                          ),
+                  (self.legends_matrix, None,                           self.highlightBondLists_matrix),
+                  (None,                self.highlightAtomLists_matrix, self.highlightBondLists_matrix),
+                  # All three other matrices: 1 parameter set
+                  (self.legends_matrix, self.highlightAtomLists_matrix, self.highlightBondLists_matrix),
+                  ]
+
+
   def test_interactive(self):
     # We avoid checking in the code with development flag set
     self.assertFalse(self.showAllImages)
@@ -213,6 +255,122 @@ def testGridSVG(self):
                                useSVG=True, drawOptions=dopts)
     self.assertIn("class='note'", svg)
 
+  def testMolsMatrixToLinear(self):
+    mols, molsPerRow, legends, highlightAtomLists, highlightBondLists = Draw._MolsNestedToLinear(
+        self.mols_matrix, self.legends_matrix, self.highlightAtomLists_matrix, self.highlightBondLists_matrix)
+
+    nrows = len(self.mols_matrix)
+
+    def _nestedOrder(self, mols_matrix, legends_matrix, highlightAtomLists_matrix, highlightBondLists_matrix):
+      for r, row in enumerate(mols_matrix):
+          for c, item in enumerate(row):
+              linear_index = (r * molsPerRow) + c
+              # Test that items in 2D list are in correct position in 1D list
+              self.assertTrue(mols[linear_index] == item)
+              # Test that 1D list items are not lists
+              self.assertFalse(isinstance(item, list))
+
+      # Other three matrices (legends; atom and bond highlighting) need not be supplied;
+      #   only test each if it's supplied
+      if legends_matrix is not None:
+          for r, row in enumerate(legends_matrix):
+              for c, item in enumerate(row):
+                  linear_index = (r * molsPerRow) + c
+                  # Test that items in 2D list are in correct position in 1D list
+                  self.assertTrue(legends[linear_index] == item)
+                  # Test that 1D list items are not lists
+                  self.assertFalse(isinstance(item, list))
+
+      if highlightAtomLists_matrix is not None:
+          for r, row in enumerate(highlightAtomLists_matrix):
+              for c, item in enumerate(row):
+                  linear_index = (r * molsPerRow) + c
+                  # Test that items in 2D list are in correct position in 1D list
+                  self.assertTrue(highlightAtomLists[linear_index] == item)
+                  # For highlight parameters, entries are lists, so check that sub-items are not lists
+                  for subitem in item:
+                      self.assertFalse(isinstance(subitem, list))
+
+      if highlightBondLists_matrix is not None:
+          for r, row in enumerate(highlightBondLists_matrix):
+              for c, item in enumerate(row):
+                  linear_index = (r * molsPerRow) + c
+                  # Test that items in 2D list are in correct position in 1D list
+                  self.assertTrue(highlightBondLists[linear_index] == item)
+                  # For highlight parameters, entries are lists, so check that sub-items are not lists
+                  for subitem in item:
+                      self.assertFalse(isinstance(subitem, list))
+
+      # Test that 1D list has the correct length
+      self.assertTrue(len(mols) == nrows * molsPerRow)
+
+    # Parametrize tests: In addition to supplying mols_matrix, supply 0-3 other matrices
+    for param_set in self.param_sets:
+      _nestedOrder(self, self.mols_matrix, *param_set)
+
+    ## Test that exceptions are thrown appropriately
+
+    # Test that supplying a non-nested list raises a ValueError
+    # Set up non-nested lists = first sublist of nested lists
+    mols_not_nested = self.mols_matrix[0]
+    legends_not_nested = self.legends_matrix[0]
+    highlightAtomLists_not_nested = self.highlightAtomLists_matrix[0]
+    highlightBondLists_not_nested = self.highlightBondLists_matrix[0]
+
+    with self.assertRaises(ValueError):
+      Draw._MolsNestedToLinear(mols_not_nested)
+
+    with self.assertRaises(ValueError):
+      Draw._MolsNestedToLinear(mols_matrix=self.mols_matrix, legends_matrix=legends_not_nested)
+
+    with self.assertRaises(ValueError):
+      Draw._MolsNestedToLinear(mols_matrix=self.mols_matrix, highlightAtomLists_matrix=highlightAtomLists_not_nested)
+
+    with self.assertRaises(ValueError):
+      Draw._MolsNestedToLinear(mols_matrix=self.mols_matrix, highlightBondLists_matrix=highlightBondLists_not_nested)
+
+    # Test that raises ValueError if other matrices aren't same size (# rows) as mols_matrix
+    with self.assertRaises(ValueError):
+      Draw._MolsNestedToLinear(mols_matrix=self.mols_matrix, legends_matrix=self.legends_matrix[0:1])
+
+    with self.assertRaises(ValueError):
+      Draw._MolsNestedToLinear(mols_matrix=self.mols_matrix, highlightAtomLists_matrix=self.highlightAtomLists_matrix[0:1])
+
+    with self.assertRaises(ValueError):
+      Draw._MolsNestedToLinear(mols_matrix=self.mols_matrix, highlightBondLists_matrix=self.highlightBondLists_matrix[0:1])
+
+    # Test that raises ValueError if other matrices' rows aren't same length as mols_matrix's corresponding row
+    # Remove last element from first row of each other matrix
+    legends_matrix_short_row0 = [self.legends_matrix[0][0:-1]] + [self.legends_matrix[1:]]
+    highlightAtomLists_matrix_short_row0 = [self.highlightAtomLists_matrix[0][0:-1]] + [self.highlightAtomLists_matrix[1:]]
+    highlightBondLists_matrix_short_row0 = [self.highlightBondLists_matrix[0][0:-1]] + [self.highlightBondLists_matrix[1:]]
+
+    with self.assertRaises(ValueError):
+      Draw._MolsNestedToLinear(mols_matrix=self.mols_matrix, legends_matrix=legends_matrix_short_row0)
+
+    with self.assertRaises(ValueError):
+      Draw._MolsNestedToLinear(mols_matrix=self.mols_matrix, highlightAtomLists_matrix=highlightAtomLists_matrix_short_row0)
+
+    with self.assertRaises(ValueError):
+      Draw._MolsNestedToLinear(mols_matrix=self.mols_matrix, highlightBondLists_matrix=highlightBondLists_matrix_short_row0)
+
+
+  def testMolsMatrixToGridImage(self):
+    # Tests are that running Draw.MolsMatrixToGridImage doesn't give an error
+    # Parametrize tests: In addition to supplying mols_matrix, supply 0-3 other matrices
+    subImgSize = (200, 200)
+
+    kwargs_value = Draw.rdMolDraw2D.MolDrawOptions()
+    kwargs_value.addAtomIndices = True
+
+    for (legends_matrix, highlightAtomLists_matrix, highlightBondLists_matrix) in self.param_sets:
+      for useSVG in (True, False):
+        for returnPNG in (True, False):
+          dwg_subImgSize_nokwargs = Draw.MolsMatrixToGridImage(self.mols_matrix, subImgSize, legends_matrix=legends_matrix, highlightAtomLists_matrix=highlightAtomLists_matrix, highlightBondLists_matrix=highlightBondLists_matrix, useSVG=useSVG, returnPNG=returnPNG)
+          dwg_subImgSize_kwargs = Draw.MolsMatrixToGridImage(self.mols_matrix, subImgSize, legends_matrix=legends_matrix, highlightAtomLists_matrix=highlightAtomLists_matrix, highlightBondLists_matrix=highlightBondLists_matrix, useSVG=useSVG, returnPNG=returnPNG, drawOptions=kwargs_value)
+          dwg_nosubImgSize_nokwargs = Draw.MolsMatrixToGridImage(self.mols_matrix, legends_matrix=legends_matrix, highlightAtomLists_matrix=highlightAtomLists_matrix, highlightBondLists_matrix=highlightBondLists_matrix, useSVG=useSVG, returnPNG=returnPNG)
+          dwg_nosubImgSize_kwargs = Draw.MolsMatrixToGridImage(self.mols_matrix, legends_matrix=legends_matrix, highlightAtomLists_matrix=highlightAtomLists_matrix, highlightBondLists_matrix=highlightBondLists_matrix, useSVG=useSVG, returnPNG=returnPNG, drawOptions=kwargs_value)
+
   def testDrawMorgan(self):
     m = Chem.MolFromSmiles('c1ccccc1CC1CC1')
     bi = {}