Repository: cbouy/mols2grid (language: python)

Issue:
Bug & feature request with selections
### Discussed in https://github.com/cbouy/mols2grid/discussions/19

<div type='discussions-op-text'>

<sup>Originally posted by **PatWalters** October 17, 2021</sup>
The callback feature is very useful, thanks!  I've been trying to figure out how to implement an enhanced cluster viewer.   I've been able to put together an example where I can click on a cluster center in one grid and show the cluster members in a second grid.  The part I'm missing is capturing the selections in the cluster member grid.  If I click on a cluster in the cluster center grid a second time, the cluster member selections are cleared.  Is there a way around this?   My example code is in this gist.  
https://gist.github.com/PatWalters/ffffa4612e2c143ab7389e539a5c88f7

Also, it would be useful to be able to have a callback in the cluster center grid when checkboxes are not present.  However, it appears that when selection=False, clicks in the grid are not recognized. </div>

* **Bug**: no callback when `selection=False`
* **Feature request**: restore checkbox state when re-displaying a grid


Discussion / hints:

Fix patch:
diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
index fba7c42..7edcffc 100644
--- a/.github/workflows/build.yml
+++ b/.github/workflows/build.yml
@@ -18,7 +18,7 @@ jobs:
     - name: Install python with pip
       uses: actions/setup-python@v2
       with:
-        python-version: 3.6
+        python-version: 3.7
         architecture: x64
     
     - name: Install dependencies for packaging
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 72ff6ba..a7a44fd 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -23,17 +23,27 @@ jobs:
     strategy:
       matrix:
         os: [ubuntu-latest]
-        python-version: [3.6, 3.8]
+        python-version: [3.8, 3.9]
         rdkit-version: ["rdkit>2020.03.1"]
         include:
         - name: CI min version
           os: ubuntu-latest
-          python-version: 3.6
+          python-version: 3.7
           rdkit-version: "rdkit==2020.03.1"
 
     steps:
       - uses: actions/checkout@v2
 
+      - name: Install Firefox
+        uses: browser-actions/setup-firefox@latest
+      
+      - run: firefox --version
+
+      - name: Prepare Selenium
+        uses: browser-actions/setup-geckodriver@latest
+      
+      - run: geckodriver --version
+
       - name: Cache conda
         uses: actions/cache@v2
         env:
@@ -44,6 +54,14 @@ jobs:
             ${{ runner.os }}-conda-${{ env.CACHE_NUMBER }}-${{
             hashFiles('etc/example-environment.yml') }}
 
+      - name: Cache pip
+        uses: actions/cache@v2
+        with:
+          path: ~/.cache/pip
+          key: ${{ runner.os }}-pip-${{ hashFiles('**/requirements.txt') }}
+          restore-keys: |
+            ${{ runner.os }}-pip-
+
       - name: Setup Conda
         uses: conda-incubator/setup-miniconda@v2
         with:
@@ -72,7 +90,23 @@ jobs:
       
       - name: Run tests
         run: |
-          pytest --color=yes --disable-pytest-warnings --cov=mols2grid --cov-report=xml tests/
+          pytest --color=yes --disable-pytest-warnings \
+                 --cov=mols2grid --cov-report=xml \
+                 tests/ -m "not webdriver"
+
+      - name: Run webdriver tests
+        uses: nick-invision/retry@v2
+        with:
+          timeout_minutes: 5
+          max_attempts: 5
+          retry_on: timeout
+          shell: bash
+          command: |
+            source /usr/share/miniconda/etc/profile.d/conda.sh
+            conda activate test
+            pytest --color=yes --disable-pytest-warnings \
+                   --cov=mols2grid --cov-report=xml --cov-append \
+                   tests/ -m "webdriver"
 
       - name: Measure tests coverage
         uses: codecov/codecov-action@v1
diff --git a/CHANGELOG.md b/CHANGELOG.md
index d37f234..94ee780 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -4,6 +4,43 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [0.2.0] - 2022/02/10
+### Added
+- `cache_selection=True` allows to retrieve the checkbox state when re-displaying a grid,
+  as long as they have the same name. Fixes #22
+- `prerender=False` moves the rendering of molecule images from Python to the browser and
+  only when the molecule is on the current page, giving a performance boost and allowing 
+  to process much larger files. Fixes #17
+- `substruct_highlight=True` highlight the atoms that matched the substructure query when
+  using the SMARTS search (only available when `prerender=False`). Fixes #18
+- Added CSV save option. Exports all the data present in `subset` and `tooltip` for the
+  current selection
+- Support for `.sdf.gz` files
+- Added automated tests of the interface, which should prevent future updates from
+  breaking things
+### Changed
+- Python 3.6 is no longer supported
+- Molecule images are now generated by the web browser (see `prerender=False` argument)
+- The coordinates of the input file are now ignored by default (`use_coords=False`). This
+  change was made to comply with generating images from SMILES string with the browser by
+  default.
+- Python callbacks are now automatically registered in Google Colab
+- Javascript callbacks can access RDKit as either `RDKit` or `RDKitModule`
+- The "img" field is now available from the callback data
+- The `subset` parameter now throws an error if "img" is not present
+- Clicking "Check all"/"Uncheck all" should now be faster
+- Bumped RDKit JS version to `2021.9.4` to better support moldrawoptions
+- Installation now requires `jinja2>=2.11.0` to prevent an error when given a pathlib.Path
+  object instead of a string
+### Fixed
+- Callbacks now work when `selection=False`. Fixes: Issue #22
+- Using both `transform` and `style` should now display the labels as expected in the
+  tooltip
+- Fixed a race condition when clicking checkboxes on different grids
+- Fixed the `gap` argument not being properly taken into account
+- Automatic resizing of the iframe (used in `mols2Grid.display`) should now work even
+  better
+
 ## [0.1.0] - 2021/10/11
 ### Added
 - The grid can be filtered using pandas DataFrame's `query` and `loc` logic (mostly
diff --git a/README.md b/README.md
index e07e4d7..b9c8020 100644
--- a/README.md
+++ b/README.md
@@ -14,12 +14,12 @@ mols2grid is an interactive chemical viewer for 2D structures of small molecules
 
 ![Demo showing mols2grid's integration in a Jupyter notebook](https://raw.githubusercontent.com/cbouy/mols2grid/master/demo.png)
 
-➡️ Try the demo notebook on [Google Colab](https://colab.research.google.com/github/cbouy/mols2grid/blob/master/demo.ipynb)
+➡️ Try the [demo notebook](https://colab.research.google.com/github/cbouy/mols2grid/blob/master/demo.ipynb) on Google Colab, or the more [advanced notebook](https://colab.research.google.com/github/rdkit/UGM_2021/blob/main/Notebooks/Bouysset_mols2grid.ipynb) from the RDKit UGM 2021.
 
 ## 🐍 Installation
 ---
 
-mols2grid was developped for Python 3.6+ and requires rdkit (>=2020.03.1), pandas and jinja2 as dependencies.  
+mols2grid was developped for Python 3.7+ and requires rdkit (>=2020.03.1), pandas and jinja2 as dependencies.  
 The easiest way to install it is from conda:
 ```shell
 conda install -c conda-forge mols2grid
@@ -30,7 +30,7 @@ Alternatively, you can also use pip:
 pip install mols2grid
 ```
 
-It is compatible with Jupyter Notebook and Google Colab (Visual Studio notebooks and Jupyterlab are not supported) and can run on Streamlit.
+It is compatible with Jupyter Notebook and Google Colab (Visual Studio notebooks and Jupyterlab are not fully supported) and can run on Streamlit.
 
 ## 📜 Usage
 ---
@@ -60,7 +60,7 @@ mols2grid.display("path/to/molecules.sdf",
 You can setup the grid from various inputs:
 * a pandas **DataFrame** (with a column of SMILES or RDKit molecules, controlled by the `smiles_col` and `mol_col` parameters),
 * a list of **RDKit molecules** (with properties accessible through the `mol.GetPropsAsDict()` method),
-* or an **SDF file**
+* or an **SDF file** (`.sdf` or `.sdf.gz`)
 
 You can also rename each field of your input with the `rename` parameter. Please note that 3 fields are automatically added regardless of your input: `mols2grid-id`, `SMILES` and `img`. If a "SMILES" field already exists, it will not be overwritten.
 
@@ -70,8 +70,10 @@ You can also rename each field of your input with the `rename` parameter. Please
 * `coordGen=True`: use the coordGen library instead of the RDKit one to depict the molecules in 2D
 * `removeHs=False`: remove explicit hydrogen atoms from the drawings
 * `size=(160, 120)`: size of each image
-* `use_coords=True`: use the coordinates of the input molecules if available
-* `MolDrawOptions=None`: RDKit's [MolDrawOptions](https://www.rdkit.org/docs/source/rdkit.Chem.Draw.rdMolDraw2D.html#rdkit.Chem.Draw.rdMolDraw2D.MolDrawOptions) class. Useful for making highly customized drawings. You can also leave this to `None`, and directly use the attributes of this class as parameters like `addStereoAnnotation=True`
+* `use_coords=False`: use the coordinates of the input molecules if available
+* `MolDrawOptions=None`: RDKit's [MolDrawOptions](https://www.rdkit.org/docs/source/rdkit.Chem.Draw.rdMolDraw2D.html#rdkit.Chem.Draw.rdMolDraw2D.MolDrawOptions) class. Useful for making highly customized drawings. You can also leave this to `None`, and directly use the attributes of this class as parameters like `addStereoAnnotation=True`. Additionally, `atomColourPalette` is made available to modify the color palette if you are not prerendering images.
+* `prerender=False`: Prerender images for the entire dataset, or generate them on-the-fly when needed
+* `cache_selection=False` : Restores the selection from a previous grid with the same name.
 
 #### Parameters for the grid
   
@@ -98,8 +100,8 @@ For the `pages` template, the following parameters are available:
     Number of rows per page
 * `border="1px solid #cccccc"`: str  
     Styling of the border around each cell (CSS)
-* `gap=0`: int or str  
-    Size of the margin around each cell (CSS)
+* `gap=0`: int
+    Size of the margin around each cell in px
 * `fontsize="12pt"`: str  
     Font size of the text displayed in each cell (CSS)
 * `fontfamily"'DejaVu', sans-serif"`: str  
@@ -127,14 +129,18 @@ For the `pages` template, the following parameters are available:
     These transformations only affect columns in `subset` and `tooltip` and do not interfere with `style`.
 * `selection=True` : bool  
     Enables the selection of molecules using a checkbox. Only usefull in the context of a Jupyter notebook. You can retrieve your selection of molecules (index and SMILES) through `mols2grid.get_selection()`
+* `cache_selection=False` : bool
+    Restores the selection from a previous grid with the same name if `True`
 * `custom_css=None` : str or None  
     Custom CSS properties applied to the content of the HTML document
 * `custom_header=None` : str or None  
     Custom libraries (CSS or JS) to be loaded in the header of the document
 * `callback=None` : str or callable  
     JavaScript or Python callback to be executed when clicking on an image. A dictionnary containing the data for the full cell is directly available as `data` in JS. For Python, the callback function must have `data` as the first argument to the function. All the values in the `data` dict are parsed as strings, except "mols2grid-id" which is always an integer.
-* `sort_by` : str or None  
+* `sort_by=None` : str or None  
     Sort the grid according to the following field (which must be present in `subset` or `tooltip`).
+* `substruct_highlight=True` : bool
+    Highlight the query when using a SMARTS substructure search. Only available when `prerender=False`.
 
 Less options are available for the `table` template, you can check the complete list of arguments with `help(mols2grid.MolGrid.to_table)`
 
@@ -151,6 +157,7 @@ You can either:
   * [Blog post](https://blog.reverielabs.com/building-web-applications-from-python-scripts-with-streamlit/) by Justin Chavez
   * [Video tutorial](https://www.youtube.com/watch?v=0rqIwSeUImo) by Data Professor
 * [Viewing clustered chemical structures](https://practicalcheminformatics.blogspot.com/2021/07/viewing-clustered-chemical-structures.html) by Pat Walters
+* [RDKit UGM 2021 notebook](https://colab.research.google.com/github/rdkit/UGM_2021/blob/main/Notebooks/Bouysset_mols2grid.ipynb)
 
 ## 👏 Acknowledgments
 ---
diff --git a/mols2grid/_version.py b/mols2grid/_version.py
index a68927d..49f34f4 100644
--- a/mols2grid/_version.py
+++ b/mols2grid/_version.py
@@ -1,1 +1,1 @@
-__version__ = "0.1.0"
\ No newline at end of file
+__version__ = "0.2.0"
\ No newline at end of file
diff --git a/mols2grid/dispatch.py b/mols2grid/dispatch.py
index 5f067f5..7674468 100644
--- a/mols2grid/dispatch.py
+++ b/mols2grid/dispatch.py
@@ -27,7 +27,7 @@ def _prepare_kwargs(kwargs, kind):
 @singledispatch
 def display(arg, **kwargs):
     """Display molecules on an interactive grid
-    
+
     Parameters
     ----------
     arg : pandas.DataFrame, SDF file or list of molecules
@@ -41,13 +41,21 @@ def display(arg, **kwargs):
     coordGen : bool (True)
         Use the coordGen library instead of the RDKit one to depict the
         molecules in 2D
-    use_coords : bool (True)
+    use_coords : bool (False)
         Use the coordinates of the molecules (only relevant when an SDF file, a
         list of molecules or a DataFrame of RDKit molecules were used as input)
-    remove_Hs : bool (False)
+    removeHs : bool (False)
         Remove hydrogen atoms from the drawings
     size : tuple ((160, 120))
         Size of each image
+    prerender : bool (False)
+        Prerender images for the entire dataset, or generate them on-the-fly
+        when needed
+    substruct_highlight : bool
+        Highlight substructure when using the SMARTS search. Only available
+        when `prerender=False`
+    MolDrawOptions : rdkit.Chem.Draw.MolDrawOptions (None)
+        Drawing options. Useful for making highly customized drawings
     rename : dict (None)
         Rename the properties in the final document
     name : str ("default")
@@ -81,8 +89,8 @@ def display(arg, **kwargs):
         Only available for the "pages" template. Number of rows per page
     border : str ("1px solid #cccccc")
         Styling of the border around each cell (CSS)
-    gap : int or str (0)
-        Size of the margin around each cell (CSS)
+    gap : int (0)
+        Size of the margin around each cell in px
     fontsize : str ("12pt")
         Font size of the text displayed in each cell (CSS)
     fontfamily : str ("'DejaVu', sans-serif")
@@ -110,6 +118,8 @@ def display(arg, **kwargs):
         displays a checkbox at the top of each cell. To access your selection (index and
         SMILES), use `mols2grid.get_selection()` or the export options in the bottom
         checkbox dropdown menu.
+    cache_selection : bool (False)
+        Restores the selection from a previous grid with the same name
     transform : dict (None)
         Functions applied to specific items in all cells. The dict must follow a
         `key: function` structure where the key must correspond to one of the columns
@@ -121,9 +131,11 @@ def display(arg, **kwargs):
         Celsius instead of Fahrenheit with a single digit precision and some text
         before (MP) and after (°C) the value. These transformations only affect
         columns in `subset` and `tooltip`, and do not interfere with `style`.
-    custom_css : str
+    custom_css : str or None
         Only available for the "pages" template. Custom CSS properties applied to the
         content of the HTML document.
+    custom_header : str or None
+        Custom libraries to be loaded in the header of the document
     callback : str
         Only available for the "pages" template. JavaScript callback to be executed when
         clicking on an image. A dictionnary containing the data for the full cell is
@@ -132,11 +144,13 @@ def display(arg, **kwargs):
     sort_by : str or None
         Sort the grid according to the following field (which must be present in `subset`
         or `tooltip`).
-    
+
     Notes
     -----
     You can also directly use RDKit's MolDrawOptions parameters as arguments.
     For more info, see https://www.rdkit.org/docs/source/rdkit.Chem.Draw.rdMolDraw2D.html#rdkit.Chem.Draw.rdMolDraw2D.MolDrawOptions
+    Additionally, `atomColourPalette` is available to customize the atom
+    palette if you're not prerendering image (`prerender=False`).
     """
     raise TypeError(f"No display method registered for type {type(arg)!r}")
 
@@ -162,14 +176,14 @@ def _(mols, **kwargs):
 @singledispatch
 def save(arg, **kwargs):
     """Generate an interactive grid of molecules and save it
-    
+
     Parameters
     ----------
     arg : pandas.DataFrame, SDF file or list of molecules
         The input containing your molecules
     output : str
         Name and path of the output document
-    
+
     See `mols2grid.display` for the full list of arguments
     """
     raise TypeError(f"No save method registered for type {type(arg)!r}")
diff --git a/mols2grid/molgrid.py b/mols2grid/molgrid.py
index 0766837..d65fa3c 100644
--- a/mols2grid/molgrid.py
+++ b/mols2grid/molgrid.py
@@ -1,4 +1,3 @@
-from typing import Type
 import warnings
 from base64 import b64encode
 from html import escape
@@ -30,8 +29,10 @@ class MolGrid:
     """
 
     def __init__(self, df, smiles_col="SMILES", mol_col=None, removeHs=False,
-        use_coords=True, coordGen=True, useSVG=True, size=(160, 120),
-        MolDrawOptions=None, rename=None, name="default", **kwargs):
+        use_coords=False, coordGen=True, useSVG=True, size=(160, 120),
+        MolDrawOptions=None, rename=None, name="default", prerender=False,
+        cache_selection=False,
+        **kwargs):
         """
         Parameters
         ----------
@@ -62,23 +63,52 @@ def __init__(self, df, smiles_col="SMILES", mol_col=None, removeHs=False,
         name : str
             Name of the grid. Used when retrieving selections from multiple
             grids at the same time
+        cache_selection : bool
+            Restores the selection from a previous grid with the same name
+        prerender : bool
+            Prerender images for the entire dataset, or generate them
+            on-the-fly when needed
         kwargs : object
-            MolDrawOptions attributes
+            MolDrawOptions attributes, and the additional `atomColourPalette`
 
         Notes
         -----
         The list of supported MolDrawOptions attributes are available in
         https://www.rdkit.org/docs/source/rdkit.Chem.Draw.rdMolDraw2D.html#rdkit.Chem.Draw.rdMolDraw2D.MolDrawOptions
 
-        ..versionchanged: 0.1.0
+        On-the-fly rendering of images does not read the atom colour palette
+        from the MolDrawOptions passed. If this is needed, use the following::
+
+            >>> MolGrid(df, atomColourPalette={1: (.8, 0, 1)})
+
+        .. versionchanged:: 0.1.0
             Added `rename` argument to replace `mapping`
+
+        .. versionadded:: 0.2.0
+            Added `prerender` and `cache_selection` arguments
+
+        .. versionchanged:: 0.2.0
+            Images are now generated on-the-fly. `use_coords` is now `False` by
+            default to avoid a systematic error when using `MolGrid.from_sdf`
         """
         if not (smiles_col or mol_col):
             raise ValueError("One of `smiles_col` or `mol_col` must be set")
         if not isinstance(name, str):
             raise TypeError(
                 f"`name` must be a string. Currently of type {type(name).__name__}")
-        Draw.rdDepictor.SetPreferCoordGen(coordGen)
+        if not prerender:
+            if not useSVG:
+                raise ValueError("On-the-fly rendering of PNG images not supported")
+            if use_coords and mol_col:
+                raise ValueError("Cannot use coordinates with on-the-fly rendering")
+        self.prefer_coordGen = coordGen
+        self.removeHs = removeHs
+        self.useSVG = useSVG
+        self.use_coords = use_coords
+        self.img_size = size
+        self.prerender = prerender
+        self.smiles_col = smiles_col
+        self.mol_col = mol_col
         if isinstance(df, pd.DataFrame):
             dataframe = df.copy()
         else:
@@ -94,50 +124,51 @@ def __init__(self, df, smiles_col="SMILES", mol_col=None, removeHs=False,
         if rename:
             dataframe.rename(columns=rename, inplace=True)
         self._extra_columns = ["img", "mols2grid-id"]
-        # generate temporary RDKit molecules
-        if smiles_col and not mol_col:
-            mol_col = "mol"
-            keep_mols = False
-            dataframe[mol_col] = dataframe[smiles_col].apply(Chem.MolFromSmiles)
-        else:
-            keep_mols = True
-        # remove hydrogens
-        if removeHs:
-            dataframe[mol_col] = dataframe[mol_col].apply(Chem.RemoveHs)
-        if not use_coords:
-            dataframe[mol_col] = dataframe[mol_col].apply(remove_coordinates)
-        # generate smiles col
-        if mol_col and (smiles_col not in dataframe.columns):
-            dataframe[smiles_col] = dataframe[mol_col].apply(mol_to_smiles)
         # add index
         dataframe["mols2grid-id"] = list(range(len(dataframe)))
-        # drop None
-        dataframe.dropna(axis=0, subset=[mol_col], inplace=True)
-        # generate drawings
-        self.useSVG = useSVG
-        opts = MolDrawOptions or Draw.MolDrawOptions()
-        for key, value in kwargs.items():
-            setattr(opts, key, value)
-        self.MolDrawOptions = opts
-        self._MolDraw2D = Draw.MolDraw2DSVG if useSVG else Draw.MolDraw2DCairo
-        self.img_size = size
-        dataframe["img"] = dataframe[mol_col].apply(self.mol_to_img)
-        if keep_mols:
-            self.dataframe = dataframe
+        # generate drawing options
+        if prerender:
+            Draw.rdDepictor.SetPreferCoordGen(coordGen)
+            opts = MolDrawOptions or Draw.MolDrawOptions()
+            for key, value in kwargs.items():
+                setattr(opts, key, value)
+            self.MolDrawOptions = opts
+            self._MolDraw2D = Draw.MolDraw2DSVG if useSVG else Draw.MolDraw2DCairo
         else:
-            self.dataframe = dataframe.drop(columns=mol_col)
-            mol_col = None
-        self.smiles_col = smiles_col
-        self.mol_col = mol_col
+            opts = {}
+            if MolDrawOptions:
+                for key in dir(MolDrawOptions):
+                    value = getattr(MolDrawOptions, key)
+                    if not (key.startswith("_")
+                            or callable(value)
+                            or value.__class__.__module__ != "builtins"):
+                        opts[key] = value
+            opts.update(kwargs)
+            opts.update({"width": self.img_size[0],
+                         "height": self.img_size[1]})
+            self.json_draw_opts = json.dumps(opts)
+        # prepare smiles and images
+        self._prepare_dataframe(dataframe)
+        self.dataframe = dataframe
         # register instance
         self._grid_id = name
-        register._init_grid(name)
+        if cache_selection:
+            try:
+                self._cached_selection = register.get_selection(name)
+            except KeyError:
+                self._cached_selection = None
+                register._init_grid(name)
+            else:
+                register._update_current_grid(name)
+        else:
+            self._cached_selection = None
+            register._init_grid(name)
 
     @classmethod
     def from_mols(cls, mols, **kwargs):
         """Set up the dataframe used by mols2grid directly from a list of RDKit
         molecules
-        
+
         Parameters
         ----------
         mols : list
@@ -153,13 +184,16 @@ def from_mols(cls, mols, **kwargs):
     @classmethod
     def from_sdf(cls, sdf_file, **kwargs):
         """Set up the dataframe used by mols2grid directly from an SDFile
-        
+
         Parameters
         ----------
         sdf_file : str
-            Path to the SDF file
+            Path to the SDF file (.sdf or .sdf.gz)
         kwargs : object
             Other arguments passed on initialization
+        
+        .. versionchanged:: 0.2.0
+            Added support for `.sdf.gz` files
         """
         mol_col = kwargs.pop("mol_col", "mol")
         df = sdf_to_dataframe(sdf_file, mol_col=mol_col)
@@ -198,33 +232,67 @@ def mol_to_img(self, mol):
         data = b64encode(img).decode()
         return f'<img src="data:image/png;base64,{data}">'
 
+    def _prepare_dataframe(self, dataframe):
+        """Prepares the dataframe with SMILES and images depending on user input. The 
+        dataframe is modified inplace."""
+        if self.prerender:
+            if self.mol_col:
+                keep_mols = True
+            else:
+                # make temporary mol column if not present
+                self.mol_col = "mol"
+                keep_mols = False
+                dataframe[self.mol_col] = (
+                    dataframe[self.smiles_col].apply(Chem.MolFromSmiles))
+            # drop empty mols
+            dataframe.dropna(axis=0, subset=[self.mol_col], inplace=True)
+            # modify mol according to user pref
+            if not self.use_coords:
+                dataframe[self.mol_col] = (
+                    dataframe[self.mol_col].apply(remove_coordinates))
+            if self.removeHs:
+                dataframe[self.mol_col] = dataframe[self.mol_col].apply(Chem.RemoveHs)
+            # cleanup
+            if not keep_mols:
+                dataframe.drop(columns=self.mol_col, inplace=True)
+                self.mol_col = None
+        # generate smiles col if not present or needs to be updated
+        if self.mol_col and self.smiles_col not in dataframe.columns:
+            dataframe[self.smiles_col] = dataframe[self.mol_col].apply(mol_to_smiles)
+        # render
+        if self.prerender:
+            dataframe["img"] = dataframe[self.mol_col].apply(self.mol_to_img)
+        else:
+            dataframe["img"] = None
+
     def render(self, template="pages", **kwargs):
         """Returns the HTML document corresponding to the "pages" or "table"
         template. See `to_pages` and `to_table` for the list of arguments
-        
+
         Parameters
         ----------
         template : str
             Kind of grid to draw:
-                * "table" is a very simple table where all molecules are
-                  displayed on the document, the main usecase is printing to
-                  PDF or on paper.
-                * "pages" is a more interactive version that splits the
-                  original data into several pages.
+
+            * "table" is a very simple table where all molecules are
+                displayed on the document, the main usecase is printing to
+                PDF or on paper.
+            * "pages" is a more interactive version that splits the
+                original data into several pages.
         """
         self.template = template
         return getattr(self, f"to_{self.template}")(**kwargs)
-    
-    def to_pages(self, subset=None, tooltip=None,
-                 cell_width=160, n_cols=5, n_rows=3,
+
+    def to_pages(self, subset=None, tooltip=None, n_cols=5, n_rows=3,
                  border="1px solid #cccccc", gap=0,
                  fontsize="12pt", fontfamily="'DejaVu', sans-serif",
                  textalign="center", tooltip_fmt="<strong>{key}</strong>: {value}",
                  tooltip_trigger="click hover", tooltip_placement="bottom",
                  hover_color="#e7e7e7", style=None, selection=True, transform=None,
-                 custom_css=None, custom_header=None, callback=None, sort_by=None):
+                 custom_css=None, custom_header=None, callback=None, sort_by=None,
+                 substruct_highlight=True):
         """Returns the HTML document for the "pages" template
-        
+
         Parameters
         ----------
         subset : list or None
@@ -248,7 +316,7 @@ def to_pages(self, subset=None, tooltip=None,
         border : str
             Styling of the border around each cell (CSS)
         gap : int
-            Size of the margin around each cell (CSS)
+            Size of the margin around each cell in px
         fontsize : str
             Font size of the text displayed in each cell (CSS)
         fontfamily : str
@@ -296,7 +364,21 @@ def to_pages(self, subset=None, tooltip=None,
         sort_by : str or None
             Sort the grid according to the following field (which must be present in
             `subset` or `tooltip`).
+        substruct_highlight : bool
+            Highlight substructure when using the SMARTS search. Only available when
+            `prerender=False`
+
+        .. versionadded:: 0.1.0
+            Added `sort_by`, `custom_css`, `custom_header` and `callback` arguments.
+            Added the ability to style an entire cell with
+            `style={"__all__": <function>}`.
+
+        .. versionadded:: 0.2.0
+            Added `substruct_highlight` argument
         """
+        if substruct_highlight and self.prerender:
+            raise ValueError(
+                "Cannot highlight substructure search with prerendered images")
         if self.mol_col:
             df = self.dataframe.drop(columns=self.mol_col).copy()
         else:
@@ -310,6 +392,8 @@ def to_pages(self, subset=None, tooltip=None,
         if subset is None:
             subset = df.columns.tolist()
             subset = [subset.pop(subset.index("img"))] + subset
+        if "img" not in subset:
+            raise KeyError("Please add the 'img' field in the `subset` parameter")
         # define fields that are searchable and sortable
         search_cols = [f"data-{col}" for col in subset if col != "img"]
         if tooltip:
@@ -350,17 +434,19 @@ def to_pages(self, subset=None, tooltip=None,
             final_columns.append(id_name)
             subset = [id_name if x == "mols2grid-id" else x for x in subset]
         # organize data
+        temp = []
         for col in subset:
             if col == "img" and tooltip:
                 s = (f'<a tabindex="0" class="data data-{col} mols2grid-tooltip" '
-                     'data-toggle="popover" data-content="foo"></a>')  
+                     'data-toggle="popover" data-content="."></a>')  
             else:
                 if style.get(col):
                     s = f'<div class="data data-{col} style-{col}" style=""></div>'
                 else:
                     s = f'<div class="data data-{col}"></div>'
-            content.append(s)
+            temp.append(s)
             column_map[col] = f"data-{col}"
+        content = temp + content
         # add but hide SMILES div if not present
         if smiles not in (subset + tooltip):
             s = f'<div class="data data-{smiles}" style="display: none;"></div>'
@@ -390,7 +476,7 @@ def to_pages(self, subset=None, tooltip=None,
             df["mols2grid-tooltip"] = df.apply(tooltip_formatter, axis=1,
                                                args=(tooltip, tooltip_fmt, style,
                                                      transform))
-            final_columns = final_columns + ["mols2grid-tooltip"]
+            final_columns += ["mols2grid-tooltip"]
             value_names = (value_names[:-1] +
                            ", {attr: 'data-content', name: 'mols2grid-tooltip'}]")
 
@@ -399,7 +485,14 @@ def to_pages(self, subset=None, tooltip=None,
             df[col] = df[col].apply(func)
 
         if selection:
-            checkbox = '<input type="checkbox" class="position-relative float-left">'
+            if self._cached_selection:
+                df["cached_checkbox"] = False
+                df.loc[df["mols2grid-id"].isin(self._cached_selection.keys()), "cached_checkbox"] = True
+                final_columns += ["cached_checkbox"]
+                value_names = (value_names[:-1] +
+                           ", {attr: 'checked', name: 'cached_checkbox'}]")
+            checkbox = ('<input type="checkbox" '
+                        'class="position-relative float-left cached_checkbox">')
         else:
             checkbox = ""
         if whole_cell_style:
@@ -416,6 +509,13 @@ def to_pages(self, subset=None, tooltip=None,
                 raise TypeError(
                     "Lambda functions are not supported as callbacks. Please "
                     "use a regular function instead.")
+            # automatically register callback in Google Colab
+            try:
+                from google import colab
+            except:
+                pass
+            else:
+                colab.output.register_callback(callback.__name__, callback)
             callback_type = "python"
             callback = callback.__name__
         else:
@@ -425,7 +525,7 @@ def to_pages(self, subset=None, tooltip=None,
             if sort_by in (subset + tooltip):
                 sort_by = f"data-{sort_by}"
             else:
-                raise ValueError(f"{sort_by} is not an available field in "
+                raise ValueError(f"{sort_by!r} is not an available field in "
                                  "`subset` or `tooltip`")
         else:
             sort_by = "mols2grid-id"
@@ -434,12 +534,13 @@ def to_pages(self, subset=None, tooltip=None,
         template = env.get_template('pages.html')
         template_kwargs = dict(
             width = width,
+            height = self.img_size[1],
             border = border,
             textalign = textalign,
             cell_width = cell_width,
             fontfamily = fontfamily,
             fontsize = fontsize,
-            gap = gap,
+            gap = f"{gap}px",
             hover_color = hover_color,
             item = item,
             item_repr = repr(item),
@@ -451,6 +552,9 @@ def to_pages(self, subset=None, tooltip=None,
             search_cols = search_cols,
             data = json.dumps(df.to_dict("records")),
             selection = selection,
+            cached_selection = ([list(self._cached_selection.keys()),
+                                 list(self._cached_selection.values())]
+                                if self._cached_selection else False),
             smiles_col = smiles,
             sort_cols = sort_cols,
             grid_id = self._grid_id,
@@ -460,6 +564,11 @@ def to_pages(self, subset=None, tooltip=None,
             callback = callback,
             callback_type = callback_type,
             sort_by = sort_by,
+            removeHs = self.removeHs,
+            prefer_coordGen = self.prefer_coordGen,
+            onthefly = not self.prerender,
+            substruct_highlight = substruct_highlight,
+            json_draw_opts = getattr(self, "json_draw_opts", ""),
         )
         return template.render(**template_kwargs)
 
@@ -476,7 +585,7 @@ def get_selection(self):
 
     def filter(self, mask):
         """Filters the grid using a mask (boolean array)
-        
+
         Parameters
         ----------
         mask : list, pd.Series, np.ndarray
@@ -492,7 +601,7 @@ def filter_by_index(self, indices):
         # convert index to mols2grid-id
         ids = self.dataframe.loc[self.dataframe.index.isin(indices)]["mols2grid-id"]
         return self._filter_by_id(ids)
-        
+
     def _filter_by_id(self, ids):
         """Filters the grid using the values in the `mols2grid-id` column"""
         if isinstance(ids, (pd.Series, np.ndarray)):
@@ -501,7 +610,7 @@ def _filter_by_id(self, ids):
             grid_id = self._grid_id,
             ids = ids)
         return Javascript(code)
-    
+
     def to_table(self, subset=None, tooltip=None, n_cols=6,
                  cell_width=160, border="1px solid #cccccc", gap=0,
                  fontsize="12pt", fontfamily="'DejaVu', sans-serif",
@@ -509,7 +618,7 @@ def to_table(self, subset=None, tooltip=None, n_cols=6,
                  tooltip_trigger="click hover", tooltip_placement="bottom",
                  hover_color="#e7e7e7", style=None, transform=None):
         """Returns the HTML document for the "table" template
-        
+
         Parameters
         ----------
         subset : list or None
@@ -605,7 +714,7 @@ def to_table(self, subset=None, tooltip=None, n_cols=6,
             td.append("\n".join(div))
             td.append("</td>")
             tr.append("\n".join(td))
-            
+
             if (ncell % n_cols == 0) or (ncell == len(df)):
                 cell = [f'<tr class="row-{nrow}">']
                 cell.append("\n".join(tr))
diff --git a/mols2grid/select.py b/mols2grid/select.py
index b787e93..bf49cbe 100644
--- a/mols2grid/select.py
+++ b/mols2grid/select.py
@@ -1,9 +1,10 @@
 import warnings
 from collections import UserDict
+from threading import Lock
 
 class SelectionRegister:
     """Register for grid selections
-    
+
     Attributes
     ----------
     SELECTIONS : dict
@@ -13,6 +14,7 @@ class SelectionRegister:
     """
     def __init__(self):
         self.SELECTIONS = {}
+        self.lock = Lock()
 
     def _update_current_grid(self, name):
         self.current_selection = name
@@ -39,10 +41,10 @@ def add_selection(self, name, _id, smiles):
         ----------
         name : str
             Name of the grid to update
-        _id : int
-            Identifier (`mols2grid-id`) of the entry
-        smiles : str
-            SMILES of the entry
+        _id : list[int]
+            List of entry identifiers (`mols2grid-id`)
+        smiles : list[str]
+            List of SMILES
 
         Notes
         -----
@@ -50,8 +52,10 @@ def add_selection(self, name, _id, smiles):
         on the grid. It is only usefull when retrieving selections from Python
         with `mols2grid.get_selection()`.
         """
-        self._update_current_grid(name)
-        self._set_selection(_id, smiles)
+        with self.lock:
+            self._update_current_grid(name)
+            for i, s in zip(_id, smiles):
+                self._set_selection(i, s)
 
     def del_selection(self, name, _id):
         """Remove an entry from a grid
@@ -60,8 +64,8 @@ def del_selection(self, name, _id):
         ----------
         name : str
             Name of the grid to update
-        _id : int
-            Identifier (`mols2grid-id`) of the entry
+        _id : list[int]
+            List of entry identifiers (`mols2grid-id`)
 
         Notes
         -----
@@ -69,8 +73,10 @@ def del_selection(self, name, _id):
         on the grid. It is only usefull when retrieving selections from Python
         with `mols2grid.get_selection()`.
         """
-        self._update_current_grid(name)
-        self._unset_selection(_id)
+        with self.lock:
+            self._update_current_grid(name)
+            for i in _id:
+                self._unset_selection(i)
 
     def get_selection(self, name=None):
         """Returns the selection for a specific MolGrid instance
diff --git a/mols2grid/templates/css/pages.css b/mols2grid/templates/css/pages.css
index 839a83f..050ddc3 100644
--- a/mols2grid/templates/css/pages.css
+++ b/mols2grid/templates/css/pages.css
@@ -19,6 +19,10 @@
 #mols2grid .cell:hover {
     background-color: {{ hover_color }} !important;
 }
+#mols2grid .cell .data-img {
+    padding: 0;
+    margin: 0;
+}
 #mols2grid .cell img, #mols2grid .cell svg {
     max-width: 100%;
     height: auto;
diff --git a/mols2grid/templates/html/checkbox_dropdown.html b/mols2grid/templates/html/checkbox_dropdown.html
index a168144..0a24613 100644
--- a/mols2grid/templates/html/checkbox_dropdown.html
+++ b/mols2grid/templates/html/checkbox_dropdown.html
@@ -1,11 +1,12 @@
 <div id="chkbox-dropdown" class="dropdown">
     <button class="btn btn-light dropdown-toggle" type="button" id="chkboxDropdown" data-toggle="dropdown" aria-haspopup="true" aria-expanded="false">☑</button>
     <div class="dropdown-menu" aria-labelledby="checkboxDropdownMenu">
-    <button id="btn-chkbox-all" class="dropdown-item" type="button">Check all</button>
-    <button id="btn-chkbox-none" class="dropdown-item" type="button">Uncheck all</button>
-    <button id="btn-chkbox-invert" class="dropdown-item" type="button">Invert</button>
-    <div class="dropdown-divider"></div>
-    <button id="btn-chkbox-copy" class="dropdown-item" type="button">Copy to clipboard</button>
-    <button id="btn-chkbox-dlsmi" class="dropdown-item" type="button">Save SMILES</button>
+        <button id="btn-chkbox-all" class="dropdown-item" type="button">Check all</button>
+        <button id="btn-chkbox-none" class="dropdown-item" type="button">Uncheck all</button>
+        <button id="btn-chkbox-invert" class="dropdown-item" type="button">Invert</button>
+        <div class="dropdown-divider"></div>
+        <button id="btn-chkbox-copy" class="dropdown-item" type="button">Copy to clipboard</button>
+        <button id="btn-chkbox-dlsmi" class="dropdown-item" type="button">Save SMILES</button>
+        <button id="btn-chkbox-dlcsv" class="dropdown-item" type="button">Save CSV</button>
     </div>
 </div>
\ No newline at end of file
diff --git a/mols2grid/templates/html/iframe.html b/mols2grid/templates/html/iframe.html
index 49d5cc6..28a5a17 100644
--- a/mols2grid/templates/html/iframe.html
+++ b/mols2grid/templates/html/iframe.html
@@ -1,11 +1,18 @@
-<iframe class="mols2grid-iframe" frameborder="0" width={{ width }}
+<script>
+    function fit_height(obj) {
+        obj.style.height = 0;
+        var height = obj.contentDocument.body.scrollHeight + {{ padding }} + 'px';
+        obj.style.height = height;
+    }
+</script>
+<iframe class="mols2grid-iframe" frameborder="0" width="{{ width }}"
 {% if height %}
 height="{{ height }}"
 {% else %}
-height="0" 
-onload="javascript:(function(o){o.style.height=o.contentWindow.document.body.scrollHeight+{{ padding }}+'px';}(this));"
+height="200"
 {% endif %}
 {% if allow %}
 allow="{{ allow }}"
 {% endif %}
-srcdoc="{{ doc }}"></iframe>
\ No newline at end of file
+srcdoc="{{ doc }}">
+</iframe>
\ No newline at end of file
diff --git a/mols2grid/templates/html/pages_header.html b/mols2grid/templates/html/pages_header.html
index 7854c47..b208628 100644
--- a/mols2grid/templates/html/pages_header.html
+++ b/mols2grid/templates/html/pages_header.html
@@ -2,4 +2,4 @@
 <link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/bootstrap@4.6.0/dist/css/bootstrap.min.css" integrity="sha384-B0vP5xmATw1+K9KRQjQERJvTumQW0nPEzvF6L/Z6nronJ3oUOFUFpCjEUQouq2+l" crossorigin="anonymous">
 <script src="https://code.jquery.com/jquery-3.5.1.slim.min.js" integrity="sha384-DfXdz2htPH0lsSSs5nCTpuj/zy4C+OGpamoFVy38MVBnE+IbbVYUew+OrCXaRkfj" crossorigin="anonymous"></script>
 <script src="https://cdn.jsdelivr.net/npm/bootstrap@4.6.0/dist/js/bootstrap.bundle.min.js" integrity="sha384-Piv4xVNRyMGpqkS2by6br4gNJ7DXjqk09RmUpJ8jgGtD7zP9yug3goQfGII0yAns" crossorigin="anonymous"></script>
-<script src="https://unpkg.com/@rdkit/rdkit@2021.3.2/Code/MinimalLib/dist/RDKit_minimal.js"></script>
\ No newline at end of file
+<script src="https://unpkg.com/@rdkit/rdkit@2021.9.4/Code/MinimalLib/dist/RDKit_minimal.js"></script>
\ No newline at end of file
diff --git a/mols2grid/templates/js/callback.js b/mols2grid/templates/js/callback.js
index 3ccd540..dd8d2cd 100644
--- a/mols2grid/templates/js/callback.js
+++ b/mols2grid/templates/js/callback.js
@@ -3,6 +3,7 @@ listObj.on("updated", function (list) {
         var data = {}
         data["mols2grid-id"] = parseInt($(this).closest(".cell")
                                                .attr("data-mols2grid-id"));
+        data["img"] = this.innerHTML;
         $(this).siblings(".data").each(function() {
             let name = this.className.split(" ")
                                      .filter(cls => cls.startsWith("data-"))[0]
diff --git a/mols2grid/templates/js/draw_mol.js b/mols2grid/templates/js/draw_mol.js
new file mode 100644
index 0000000..338327e
--- /dev/null
+++ b/mols2grid/templates/js/draw_mol.js
@@ -0,0 +1,33 @@
+// generate images for the currently displayed molecules
+RDKit.prefer_coordgen({{ prefer_coordGen | tojson }});
+function draw_mol(smiles) {
+    var mol = RDKit.get_mol(smiles, '{"removeHs": {{ removeHs | tojson }} }');
+    var svg = "";
+    if (mol.is_valid()) {
+        var highlights = smarts_matches[smiles];
+        if (highlights) {
+            var details = Object.assign({}, draw_opts, highlights);
+            details = JSON.stringify(details);
+        } else {
+            var details = json_draw_opts;
+        }
+        svg = mol.get_svg_with_highlights(details);
+    }
+    mol.delete();
+    if (svg == "") {
+        return '<svg width="{{ cell_width }}" height="{{ height }}" xmlns="http://www.w3.org/2000/svg" version="1.1" viewBox="0 0 {{ cell_width }} {{ height }}"></svg>';
+    }
+    return svg;
+}
+listObj.on("updated", function (list) {
+    var query = $('#mols2grid #searchbar').val();
+    if (query === "") {
+        smarts_matches = {};
+    }
+    $('#mols2grid .cell').each(function() {
+        var $t = $(this);
+        var smiles = $t.children(".data-{{ smiles_col }}").first().text()
+        var svg = draw_mol(smiles);
+        $t.children(".data-img").html(svg);
+    });
+});
\ No newline at end of file
diff --git a/mols2grid/templates/js/kernel.js b/mols2grid/templates/js/kernel.js
index 686345f..f7fde83 100644
--- a/mols2grid/templates/js/kernel.js
+++ b/mols2grid/templates/js/kernel.js
@@ -5,26 +5,26 @@ if (window.parent.IPython !== undefined) {
     var kernel = window.parent.IPython.notebook.kernel;
     kernel.execute('from mols2grid.select import register as _m2g_reg')
     function add_selection(grid_id, _id, smiles) {
-        SELECTION.set(_id, smiles);
-        kernel.execute("_m2g_reg.add_selection('"+grid_id+"', "+_id+",'"+smiles+"')");
+        SELECTION.multi_set(_id, smiles);
+        kernel.execute("_m2g_reg.add_selection('"+grid_id+"', "+JSON.stringify(_id)+","+JSON.stringify(smiles)+")");
     }
     function del_selection(grid_id, _id) {
-        SELECTION.delete(_id);
-        kernel.execute("_m2g_reg.del_selection('"+grid_id+"', "+_id+")");
+        SELECTION.multi_del(_id);
+        kernel.execute("_m2g_reg.del_selection('"+grid_id+"', "+JSON.stringify(_id)+")");
     }
 } else if (window.parent.google !== undefined) {
     // Google colab
     kernel_env = "colab";
     var kernel = window.parent.google.colab.kernel;
     function add_selection(grid_id, _id, smiles) {
-        SELECTION.set(_id, smiles);
+        SELECTION.multi_set(_id, smiles);
         (async function() {
         const result = await kernel.invokeFunction('_m2g_reg.add_selection',
                                                    [grid_id, _id, smiles], {});
         })();
     }
     function del_selection(grid_id, _id) {
-        SELECTION.delete(_id);
+        SELECTION.multi_del(_id);
         (async function() {
         const result = await kernel.invokeFunction('_m2g_reg.del_selection',
                                                    [grid_id, _id], {});
@@ -32,9 +32,9 @@ if (window.parent.IPython !== undefined) {
     }
 } else {
     function add_selection(grid_id, _id, smiles) {
-        SELECTION.set(_id, smiles);
+        SELECTION.multi_set(_id, smiles);
     }
     function del_selection(grid_id, _id) {
-        SELECTION.delete(_id);
+        SELECTION.multi_del(_id);
     }
 }
\ No newline at end of file
diff --git a/mols2grid/templates/js/molstorage.js b/mols2grid/templates/js/molstorage.js
index ea199bb..4b523f1 100644
--- a/mols2grid/templates/js/molstorage.js
+++ b/mols2grid/templates/js/molstorage.js
@@ -1,4 +1,14 @@
 class MolStorage extends Map {
+    multi_set(_id, _smiles) {
+        for (let i=0; i < _id.length; i++) {
+            this.set(_id[i], _smiles[i]);
+        }
+    }
+    multi_del(_id) {
+        for (let i=0; i < _id.length; i++) {
+            this.delete(_id[i]);
+        };
+    }
     to_dict() {
         var content = "{";
         for (let [key, value] of this) {
@@ -18,5 +28,6 @@ class MolStorage extends Map {
         a.href = URL.createObjectURL(file);
         a.download = fileName;
         a.click();
+        a.remove();
     }
 }
\ No newline at end of file
diff --git a/mols2grid/templates/js/pages.js b/mols2grid/templates/js/pages.js
index 11b5ffe..74f542f 100644
--- a/mols2grid/templates/js/pages.js
+++ b/mols2grid/templates/js/pages.js
@@ -1,8 +1,3 @@
-// RDKit
-onRuntimeInitialized: initRDKitModule().then(function(instance) {
-    RDKitModule = instance;
-    console.log('RDKit version: ' + RDKitModule.version());
-});
 // list.js
 var listObj = new List('mols2grid', {
     valueNames: {{ value_names }},
@@ -22,28 +17,40 @@ if (window.parent.mols2grid_lists === undefined) {
     window.parent.mols2grid_lists = {};
 }
 window.parent.mols2grid_lists[{{ grid_id | tojson }}] = listObj;
+
 {% if selection %}
 // selection
 {% include 'js/molstorage.js' %}
 var SELECTION = new MolStorage();
+{% endif %}
+
+{% if selection or callback %}
+// kernel
 {% include 'js/kernel.js' %}
 {% endif %}
+
+{% if selection and cached_selection %}
+// restore checkbox state
+SELECTION.multi_set({{ cached_selection[0] }}, {{ cached_selection[1] }});
+listObj.on("updated", function (list) {
+    $('#mols2grid .cell input[checked="false"]').prop("checked", false);
+});
+{% endif %}
+
 // sort
 {% include 'js/sort.js' %}
-// search bar
-{% include 'js/search.js' %}
+
 {% if whole_cell_style %}
 // add style for whole cell
 listObj.on("updated", function (list) {
     $('#mols2grid div.cell').each(function() {
         var $t = $(this);
-        $t.attr({
-            style: $t.attr('data-cellstyle')
-        })
+        $t.attr({style: $t.attr('data-cellstyle')})
           .removeAttr('data-cellstyle');
     });
 });
 {% endif %}
+
 {% if tooltip %}
 // tooltips
 $.fn.tooltip.Constructor.Default.whiteList.span = ['style']
@@ -51,26 +58,55 @@ listObj.on("updated", function (list) {
     $(function () {
         // hide previous popovers
         $('#mols2grid a.page-link').click(function(e) {
-        $('.mols2grid-tooltip[data-toggle="popover"]').popover('hide')
+            $('.mols2grid-tooltip[data-toggle="popover"]').popover('hide')
         });
         // create new popover
         $('.mols2grid-tooltip[data-toggle="popover"]').popover({
-        placement: {{ tooltip_placement }},
-        trigger: {{ tooltip_trigger }},
-        html: true,
-        sanitize: false,
+            placement: {{ tooltip_placement }},
+            trigger: {{ tooltip_trigger }},
+            html: true,
+            sanitize: false,
         });
     })
 });
 {% endif %}
+
 {% if selection %}
 // selection modifyers and export options
 {% include 'js/selection_actions.js' %}
 {% endif %}
+
 {% if callback %}
 // callback
 {% include 'js/callback.js' %}
 {% endif %}
-{%if tooltip or selection or whole_cell_style %}
-listObj.update();
+
+{% if onthefly %}
+// generate images for the currently displayed molecules
+var draw_opts = {{ json_draw_opts }};
+var json_draw_opts = JSON.stringify(draw_opts);
 {% endif %}
+var smarts_matches = {};
+
+// Load RDKit
+window
+.initRDKitModule()
+.then(function(RDKit) {
+    console.log('RDKit version: ', RDKit.version());
+    window.RDKit = RDKit;
+    window.RDKitModule = RDKit;
+
+    // search bar
+    {% include 'js/search.js' %}
+
+    {% if onthefly %}
+    {% include 'js/draw_mol.js' %}
+    {% endif %}
+
+    // trigger update to activate tooltips, draw images, setup callbacks...
+    listObj.update();
+    // resize iframe to fit content
+    if (window.frameElement) {
+        window.parent.fit_height(window.frameElement);
+    }
+});
\ No newline at end of file
diff --git a/mols2grid/templates/js/search.js b/mols2grid/templates/js/search.js
index 7ff4246..225072a 100644
--- a/mols2grid/templates/js/search.js
+++ b/mols2grid/templates/js/search.js
@@ -1,14 +1,23 @@
 function SmartsSearch(query, columns) {
-    query = $('#mols2grid #searchbar').val();
-    var smiles = columns[0];
-    var qmol = RDKitModule.get_qmol(query);
+    var smiles_col = columns[0];
+    smarts_matches = {};
+    var query = $('#mols2grid #searchbar').val();
+    var qmol = RDKit.get_qmol(query);
     if (qmol.is_valid()) {
         for (var k = 0, kl = listObj.items.length; k < kl; k++) {
             var item = listObj.items[k];
-            var mol = RDKitModule.get_mol(item.values()[smiles]);
-            var results = JSON.parse(mol.get_substruct_match(qmol));
-            if (results.atoms) {
-                item.found = true;
+            var smiles = item.values()[smiles_col]
+            var mol = RDKit.get_mol(smiles);
+            if (mol.is_valid()) {
+                var results = JSON.parse(mol.get_substruct_match(qmol));
+                if (results.atoms) {
+                    item.found = true;
+                    {% if onthefly and substruct_highlight %}
+                    smarts_matches[smiles] = results;
+                    {% endif %}
+                } else {
+                    item.found = false;
+                }
             } else {
                 item.found = false;
             }
@@ -26,6 +35,7 @@ $('#mols2grid .search-btn').click(function() {
 $('#mols2grid #searchbar').on("keyup", function(e) {
     var query = e.target.value;
     if (search_type === "Text") {
+        smarts_matches = {};
         listObj.search(query, {{ search_cols }});
     } else {
         listObj.search(query, ["data-{{ smiles_col }}"], SmartsSearch);
diff --git a/mols2grid/templates/js/selection_actions.js b/mols2grid/templates/js/selection_actions.js
index 3c7f482..2530e79 100644
--- a/mols2grid/templates/js/selection_actions.js
+++ b/mols2grid/templates/js/selection_actions.js
@@ -1,46 +1,61 @@
-// add selection buttons
-$("#sort-dropdown").before(`{% include 'html/checkbox_dropdown.html' %}`);
 // check all
 $('#btn-chkbox-all').click(function (e) {
     var current_page = parseInt($("li.page-item.active > a").text());
     var n_items = {{ n_items_per_page }};
     var last_page = parseInt($("li.page-item > a").last().text());
+    var _id = [];
+    var _smiles = [];
     for (let i = 0; i < last_page; i++) {
         listObj.show(i * n_items + 1, n_items);
-        $("input:checkbox:not(:checked)").each(function() {
+        $("input:checkbox").each(function() {
             this.checked = true;
-            $(this).trigger("change");
+            _id.push(parseInt($(this).closest(".cell").attr("data-mols2grid-id")));
+            _smiles.push($($(this).siblings(".data-{{ smiles_col }}")[0]).text());
         });
     }
     listObj.show((current_page - 1) * n_items + 1, n_items);
+    add_selection({{ grid_id | tojson }}, _id, _smiles);
 });
 // uncheck all
 $('#btn-chkbox-none').click(function (e) {
     var current_page = parseInt($("li.page-item.active > a").text());
     var n_items = {{ n_items_per_page }};
     var last_page = parseInt($("li.page-item > a").last().text());
+    var _id = [];
     for (let i = 0; i < last_page; i++) {
         listObj.show(i * n_items + 1, n_items);
         $("input:checkbox:checked").each(function() {
             this.checked = false;
-            $(this).trigger("change");
+            _id.push(parseInt($(this).closest(".cell").attr("data-mols2grid-id")));
         });
     }
     listObj.show((current_page - 1) * n_items + 1, n_items);
+    del_selection({{ grid_id | tojson }}, _id);
 });
 // invert
 $('#btn-chkbox-invert').click(function (e) {
     var current_page = parseInt($("li.page-item.active > a").text());
     var n_items = {{ n_items_per_page }};
     var last_page = parseInt($("li.page-item > a").last().text());
+    var _id_add = [];
+    var _id_del = [];
+    var _smiles = [];
     for (let i = 0; i < last_page; i++) {
         listObj.show(i * n_items + 1, n_items);
         $("input:checkbox").each(function() {
-        this.checked = !this.checked;
-        $(this).trigger("change");
+            this.checked = !this.checked;
+            var _id = parseInt($(this).closest(".cell").attr("data-mols2grid-id"));
+            if (this.checked) {
+                _id_add.push(_id);
+                _smiles.push($(this).siblings(".data-{{ smiles_col }}").first().text());
+            } else {
+                _id_del.push(_id);
+            }
         });
     }
     listObj.show((current_page - 1) * n_items + 1, n_items);
+    del_selection({{ grid_id | tojson }}, _id_del);
+    add_selection({{ grid_id | tojson }}, _id_add, _smiles);
 });
 // copy to clipboard
 $("#btn-chkbox-copy").click(function(e) {
@@ -50,15 +65,41 @@ $("#btn-chkbox-copy").click(function(e) {
 $("#btn-chkbox-dlsmi").click(function(e) {
     SELECTION.download_smi("selection.smi");
 });
+// export CSV
+$("#btn-chkbox-dlcsv").click(function(e) {
+    var sep = "\t"
+    // same order as subset + tooltip
+    var columns = Array.from(listObj.items[0].elm.querySelectorAll("div.data"))
+                       .map(elm => elm.classList[1]);
+    // remove 'data-'
+    var header = columns.map(name => name.slice(5));
+    // csv content
+    header = ["index"].concat(header).join(sep);
+    var content = header + "\n";
+    for (let [index, smiles] of SELECTION.entries()) {
+        var data = listObj.items[index].values();
+        content += index;
+        columns.forEach((key) => {
+            content += sep + data[key];
+        })
+        content += "\n";
+    }
+    var a = document.createElement("a");
+    var file = new Blob([content], {type: "text/csv"});
+    a.href = URL.createObjectURL(file);
+    a.download = "selection.csv";
+    a.click();
+    a.remove();
+});
 // update selection on checkbox click
 listObj.on("updated", function (list) {
     $("input:checkbox").change(function() {
         var _id = parseInt($(this).closest(".cell").attr("data-mols2grid-id"));
-        var _smiles = $($(this).siblings(".data-{{ smiles_col }}")[0]).text();
         if (this.checked) {
-            add_selection({{ grid_id | tojson }}, _id, _smiles);
+            var _smiles = $($(this).siblings(".data-{{ smiles_col }}")[0]).text();
+            add_selection({{ grid_id | tojson }}, [_id], [_smiles]);
         } else {
-            del_selection({{ grid_id | tojson }}, _id);
+            del_selection({{ grid_id | tojson }}, [_id]);
         }
     }); 
 });
\ No newline at end of file
diff --git a/mols2grid/templates/pages.html b/mols2grid/templates/pages.html
index 2aeff9d..830531a 100644
--- a/mols2grid/templates/pages.html
+++ b/mols2grid/templates/pages.html
@@ -18,6 +18,9 @@
       <div class="d-inline">
         <ul class="pagination float-left"></ul>
         <div id="controls" class="d-flex flex-row float-right">
+          {% if selection %}
+          {% include 'html/checkbox_dropdown.html' %}
+          {% endif %}
           <div id="sort-dropdown" class="dropdown pl-2">
             <button class="btn btn-light dropdown-toggle" type="button" id="sortDropdown" data-toggle="dropdown" aria-haspopup="true" aria-expanded="false">
               Sort by
@@ -30,7 +33,7 @@
               {% if loop.first %}
               <button class="dropdown-item sort-btn {{ is_active }}" type="button" data-name="mols2grid-id">Index</button>
               {% else %}
-              <button class="dropdown-item sort-btn {{ is_active }}" type="button" data-name="{{ name }}">{{ col_name | capitalize }}</button>
+              <button class="dropdown-item sort-btn {{ is_active }}" type="button" data-name="{{ name }}">{{ col_name }}</button>
               {% endif %}
               {% endfor %}
               {% if selection %}
diff --git a/mols2grid/utils.py b/mols2grid/utils.py
index cf8a9dd..a7e21d8 100644
--- a/mols2grid/utils.py
+++ b/mols2grid/utils.py
@@ -1,4 +1,4 @@
-from functools import wraps
+from functools import wraps, partial
 from importlib.util import find_spec
 from jinja2 import Environment, FileSystemLoader
 from pathlib import Path
@@ -23,7 +23,7 @@ def wrapper(*args, **kwargs):
 
 def tooltip_formatter(s, subset, fmt, style, transform):
     """Function to generate tooltips from a pandas Series
-    
+
     Parameters
     ----------
     s : pandas.Series
@@ -39,8 +39,8 @@ def tooltip_formatter(s, subset, fmt, style, transform):
     """
     items = []
     for k, v in s[subset].to_dict().items():
-        v = transform[k](v) if transform.get(k) else v
-        v = f'<span style="{style[k](v)}">{v}</span>' if style.get(k) else v
+        displayed = transform[k](v) if transform.get(k) else v
+        v = f'<span style="{style[k](v)}">{displayed}</span>' if style.get(k) else displayed
         items.append(fmt.format(key=k, value=v))
     return "<br>".join(items)
 
@@ -50,14 +50,19 @@ def mol_to_smiles(mol):
 
 def mol_to_record(mol, mol_col="mol"):
     """Function to create a dict of data from an RDKit molecule"""
-    return {"SMILES": Chem.MolToSmiles(mol),
-            **mol.GetPropsAsDict(includePrivate=True),
+    return {**mol.GetPropsAsDict(includePrivate=True),
             mol_col: mol} if mol else {}
 
 def sdf_to_dataframe(sdf_path, mol_col="mol"):
-    """Returns a dataframe of molecules from an SDF file"""
-    return pd.DataFrame([mol_to_record(mol, mol_col)
-                         for mol in Chem.SDMolSupplier(sdf_path)])
+    """Returns a dataframe of molecules from an SDfile (.sdf or .sdf.gz)"""
+    if sdf_path.endswith(".gz"):
+        import gzip
+        read_file = gzip.open
+    else:
+        read_file = partial(open, mode="rb")
+    with read_file(sdf_path) as f:
+        return pd.DataFrame([mol_to_record(mol, mol_col)
+                             for mol in Chem.ForwardSDMolSupplier(f)])
 
 def remove_coordinates(mol):
     """Removes the existing coordinates from the molecule. The molecule is
@@ -67,7 +72,7 @@ def remove_coordinates(mol):
 
 def make_popup_callback(title, html, js="", style=""):
     """Creates a JavaScript callback that displays a popup window
-    
+
     Parameters
     ----------
     title : str
diff --git a/setup.cfg b/setup.cfg
index cf93334..08935ce 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -33,12 +33,18 @@ packages = mols2grid
 python_requires = >=3.6
 install_requires =
     pandas
-    jinja2
+    jinja2>=2.11.0
 zip_safe = False
 include_package_data = True
 
 [options.extras_require]
 tests =
-    pytest>=6.2.5
-    pytest-cov>=2.12.1
-    codecov>=2.1.12
+    ipython==7.12.0
+    selenium==3.141.0
+    cairosvg==2.5.2
+    geckodriver-autoinstaller==0.1.0
+    pytest==6.2.5
+    pytest-cov==2.12.1
+    codecov==2.1.12
+    flaky==3.7.0
+    pytest-order==1.0.1


Test patch:
diff --git a/tests/pytest.ini b/tests/pytest.ini
new file mode 100644
index 0000000..3a19564
--- /dev/null
+++ b/tests/pytest.ini
@@ -0,0 +1,3 @@
+[pytest]
+markers =
+    webdriver: marks tests requiring a webdriver
\ No newline at end of file
diff --git a/tests/test_interface.py b/tests/test_interface.py
new file mode 100644
index 0000000..9450e7c
--- /dev/null
+++ b/tests/test_interface.py
@@ -0,0 +1,601 @@
+import pytest
+from flaky import flaky
+import os
+from base64 import b64encode, b64decode
+from selenium import webdriver
+from selenium.webdriver.support.ui import WebDriverWait
+from selenium.webdriver.support import expected_conditions as EC
+from selenium.webdriver.common.by import By
+from selenium.webdriver.common.keys import Keys
+from selenium.webdriver.common.action_chains import ActionChains
+from selenium.common.exceptions import (NoSuchElementException,
+                                        StaleElementReferenceException)
+import geckodriver_autoinstaller
+from rdkit import __version__ as rdkit_version
+from rdkit import Chem, RDConfig
+from rdkit.Chem import AllChem
+from PIL import Image
+from io import BytesIO
+from hashlib import md5
+from cairosvg import svg2png
+import mols2grid
+from mols2grid.utils import env
+from mols2grid.select import register
+
+geckodriver_autoinstaller.install()
+pytestmark = pytest.mark.webdriver
+
+class FirefoxDriver(webdriver.Firefox):
+    def wait_for_img_load(self, max_delay=15, selector="#mols2grid .cell .data-img svg"):
+        return (WebDriverWait(self, max_delay)
+                .until(EC.presence_of_element_located(
+                        (By.CSS_SELECTOR, selector))))
+
+    def wait(self, condition, max_delay=5):
+        return (WebDriverWait(self, max_delay,
+                      ignored_exceptions=[StaleElementReferenceException])
+                .until(condition))
+
+    def find_by_id(self, element_id, **kwargs):
+        condition = EC.presence_of_element_located((By.ID, element_id))
+        return self.wait(condition, **kwargs)
+
+    def find_by_css_selector(self, css_selector, **kwargs):
+        condition = EC.presence_of_element_located((By.CSS_SELECTOR, css_selector))
+        return self.wait(condition, **kwargs)
+
+    def find_by_class_name(self, name, **kwargs):
+        condition = EC.presence_of_element_located((By.CLASS_NAME, name))
+        return self.wait(condition, **kwargs)
+
+    def find_all_by_class_name(self, name, **kwargs):
+        condition = EC.presence_of_all_elements_located((By.CLASS_NAME, name))
+        return self.wait(condition, **kwargs)
+
+    def get_svg_md5_hash(self, selector="#mols2grid .cell .data-img"):
+        img = self.find_by_css_selector(selector)
+        im = svg2png(bytestring=(img.get_attribute("innerHTML")))
+        im = Image.open(BytesIO(im))
+        return md5(im.tobytes()).hexdigest()
+
+    def get_png_md5_hash(self, selector="#mols2grid .cell .data-img *"):
+        img = self.find_by_css_selector(selector)
+        im = Image.open(BytesIO(b64decode(img.get_attribute("src")[22:])))
+        return md5(im.tobytes()).hexdigest()
+
+
+def determine_scope(fixture_name, config):
+    if os.environ.get("GITHUB_ACTIONS", False):
+        return "function"
+    return "module"
+
+@pytest.fixture(scope=determine_scope)
+def driver():
+    options = webdriver.FirefoxOptions()
+    options.headless = True
+    driver = FirefoxDriver(options=options)
+    driver.set_page_load_timeout(10)
+    yield driver
+    driver.quit()
+
+sdf_path = f"{RDConfig.RDDocsDir}/Book/data/solubility.test.sdf"
+
+def get_grid(df, **kwargs):
+    kwargs.setdefault("mol_col", "mol")
+    return mols2grid.MolGrid(df, **kwargs)
+
+def get_doc(grid, kwargs):
+    html = grid.render(**kwargs)
+    html = b64encode(html.encode()).decode()
+    return "data:text/html;base64,{}".format(html)
+
+@pytest.fixture(scope="module")
+def df():
+    return mols2grid.sdf_to_dataframe(sdf_path)[:30]
+
+@pytest.fixture(scope="module")
+def mols(df):
+    return df["mol"][:5]
+
+@pytest.fixture(scope="module")
+def grid(df):
+    return get_grid(df)
+
+@pytest.fixture(scope="module")
+def html_doc(grid):
+    return get_doc(grid, dict(n_rows=1, subset=["_Name", "img"]))
+
+# make sure non-parametrized test is ran first
+@pytest.mark.order(1)
+def test_no_subset_all_visible(driver, grid):
+    doc = get_doc(grid, {"selection": False})
+    driver.get(doc)
+    columns = set(grid.dataframe.columns.drop("mol").to_list())
+    cell = driver.find_by_css_selector("#mols2grid .cell")
+    data_el = cell.find_elements_by_class_name("data")
+    classes = [c.replace("data-", "").replace("-copy", "")
+               for x in data_el
+               for c in x.get_attribute("class").split(" ")
+               if c.startswith("data-")]
+    classes = set(classes)
+    assert classes == columns
+
+def test_smiles_hidden(driver, html_doc):
+    driver.get(html_doc)
+    el = driver.find_by_css_selector("#mols2grid .cell .data-SMILES")
+    assert not el.is_displayed()
+
+@pytest.mark.parametrize("page", [1, 2, 3])
+@pytest.mark.parametrize("n_cols", [1, 3])
+@pytest.mark.parametrize("n_rows", [1, 3])
+def test_page_click(driver, grid, page, n_cols, n_rows):
+    doc = get_doc(grid, dict(n_cols=n_cols, n_rows=n_rows))
+    driver.get(doc)
+    for i in range(2, page+1):
+        next_page = driver.find_by_css_selector(f'a.page-link[data-i="{i}"]')
+        next_page.click()
+    first_cell = driver.find_by_class_name('cell')
+    mols2grid_id = n_cols * n_rows * (page - 1)
+    name = first_cell.find_element_by_class_name('data-_Name')
+    ref = grid.dataframe.iloc[mols2grid_id]
+    assert name.text == ref["_Name"]
+
+@pytest.mark.parametrize(["name", "css_prop", "value", "expected"], [
+    ("gap", "margin-top", 20, "20px"),
+    ("border", "border-top-width", "3px solid", "3px"),
+    ("border", "border-top-style", "1px dashed", "dashed"),
+    ("border", "border-top-color", "1px solid blue", "rgb(0, 0, 255)"),
+    ("fontsize", "font-size", "16pt", "21.3333px"),
+    ("fontfamily", "font-family", "Consolas", "Consolas"),
+    ("textalign", "text-align", "right", "right"),
+    ("custom_css", "background-color", ".cell { background-color: black; }", "rgb(0, 0, 0)"),
+])
+def test_css_properties(driver, grid, name, css_prop, value, expected):
+    doc = get_doc(grid, {name: value})
+    driver.get(doc)
+    computed = driver.execute_script(f"return getComputedStyle(document.querySelector('#mols2grid .cell')).getPropertyValue({css_prop!r});")
+    assert computed == expected
+
+def test_text_search(driver, html_doc):
+    driver.get(html_doc)
+    driver.wait_for_img_load()
+    (ActionChains(driver)
+        .send_keys_to_element(driver.find_by_id("searchbar"), "iodopropane")
+        .key_up(Keys.TAB)
+        .pause(.5)
+        .perform())
+    el = driver.find_by_css_selector("#mols2grid .cell .data-SMILES")
+    assert el.get_attribute("innerHTML") == "CC(I)C"
+
+@flaky(max_runs=3, min_passes=1)
+def test_smarts_search(driver, html_doc):
+    driver.get(html_doc)
+    driver.wait_for_img_load()
+    (ActionChains(driver)
+        .click(driver.find_by_id("searchBtn"))
+        .pause(.5)
+        .click(driver.find_by_id("smartsSearch"))
+        .send_keys_to_element(driver.find_by_id("searchbar"), "CC(I)C")
+        .key_up(Keys.TAB)
+        .pause(.5)
+        .perform())
+    el = driver.find_by_css_selector("#mols2grid .cell .data-_Name")
+    assert el.text == "2-iodopropane"
+
+def test_selection_click(driver, html_doc):
+    driver.get(html_doc)
+    driver.wait_for_img_load()
+    (ActionChains(driver)
+        .click(driver.find_by_css_selector("input[type='checkbox']"))
+        .pause(.5)
+        .perform())
+    sel = driver.execute_script("return SELECTION.to_dict();")
+    sel = eval(sel)
+    assert sel == {0: "CCC(C)CC"}
+    register._clear()
+
+def test_selection_with_cache_check_and_uncheck(driver, df):
+    register._init_grid("cached_sel")
+    register._set_selection(0, "CCC(C)CC")
+    grid = get_grid(df, name="cached_sel", cache_selection=True)
+    doc = get_doc(grid, {})
+    driver.get(doc)
+    driver.wait_for_img_load()
+    sel = driver.execute_script("return SELECTION.to_dict();")
+    sel = eval(sel)
+    assert sel == {0: "CCC(C)CC"}
+    (ActionChains(driver)
+        .click(driver.find_by_css_selector("input[type='checkbox']"))
+        .pause(.5)
+        .perform())
+    sel = driver.execute_script("return SELECTION.to_dict();")
+    sel = eval(sel)
+    assert sel == {}
+    register._clear()
+
+def test_selection_check_uncheck_invert(driver, html_doc):
+    driver.get(html_doc)
+    driver.wait_for_img_load()
+    (ActionChains(driver)
+        .click(driver.find_by_id("chkboxDropdown"))
+        .pause(.5)
+        .click(driver.find_by_id("btn-chkbox-all"))
+        .perform())
+    sel = driver.execute_script("return SELECTION.to_dict();")
+    sel = eval(sel)
+    assert len(sel) == 30
+    (ActionChains(driver)
+        .click(driver.find_by_id("chkboxDropdown"))
+        .pause(.5)
+        .click(driver.find_by_id("btn-chkbox-none"))
+        .perform())
+    sel = driver.execute_script("return SELECTION.to_dict();")
+    sel = eval(sel)
+    assert sel == {}
+    (ActionChains(driver)
+        .click(driver.find_by_css_selector("input[type='checkbox']"))
+        .click(driver.find_by_id("chkboxDropdown"))
+        .pause(.5)
+        .click(driver.find_by_id("btn-chkbox-invert"))
+        .perform())
+    sel = driver.execute_script("return SELECTION.to_dict();")
+    sel = eval(sel)
+    assert len(sel) == 29
+    register._clear()
+
+def test_check_all_selects_only_not_hidden(driver, html_doc):
+    driver.get(html_doc)
+    driver.wait_for_img_load()
+    (ActionChains(driver)
+        .send_keys_to_element(driver.find_by_id("searchbar"), "iodopropane")
+        .key_up(Keys.TAB)
+        .click(driver.find_by_id("chkboxDropdown"))
+        .pause(.3)
+        .click(driver.find_by_id("btn-chkbox-all"))
+        .perform())
+    driver.find_by_id("searchbar").clear()
+    (ActionChains(driver)
+        .send_keys_to_element(driver.find_by_id("searchbar"), Keys.BACKSPACE)
+        .key_up(Keys.TAB)
+        .perform())
+    sel = driver.execute_script("return SELECTION.to_dict();")
+    sel = eval(sel)
+    assert sel == {27: 'CC(I)C'}
+    register._clear()
+
+@pytest.mark.parametrize("prerender", [True, False])
+def test_image_size(driver, df, prerender):
+    grid = get_grid(df, size=(100, 100), prerender=prerender)
+    doc = get_doc(grid, {"selection": False, "border": "0", "substruct_highlight": False})
+    driver.get(doc)
+    if not prerender:
+        driver.wait_for_img_load()
+    img = driver.find_by_css_selector("#mols2grid .cell .data-img *")
+    assert img.size == {"height": 100.0, "width": 100.0}
+
+def test_image_use_coords(driver, df):
+    mols = df["mol"][:1]
+    AllChem.EmbedMolecule(mols[0], randomSeed=0xf00d)
+    grid = mols2grid.MolGrid.from_mols(
+        mols, use_coords=True, prerender=True, useSVG=False)
+    doc = get_doc(grid, {"substruct_highlight": False})
+    driver.get(doc)
+    md5_hash = driver.get_png_md5_hash()
+    if rdkit_version == "2020.03.1":
+        assert md5_hash == "aed60ed28347831d24f02dbb5be19007"
+    else:
+        assert md5_hash == "6909ba43f86003cea9b0fd8d723cddfe"
+
+@pytest.mark.parametrize(["coordGen", "prerender", "expected"], [
+    (True,  True, "acf7cf7cd5cfaa5eaf4a4f257e290e49"),
+    (False, True, "34a44eef3c0d2af6503069fc068e705d"),
+    (True, False, "d8227a9948b6f801193085e5a8b257a2"),
+    (False, False, "a43327f2cfcff45a2aa87217cb658eb8"),
+])
+def test_coordgen(driver, mols, coordGen, prerender, expected):
+    useSVG = not prerender
+    grid = mols2grid.MolGrid.from_mols(mols,
+        coordGen=coordGen, prerender=prerender, useSVG=useSVG, use_coords=False)
+    doc = get_doc(grid, {"substruct_highlight": False})
+    driver.get(doc)
+    if not prerender:
+        driver.wait_for_img_load()
+    if useSVG:
+        md5_hash = driver.get_svg_md5_hash()
+    else:
+        md5_hash = driver.get_png_md5_hash()
+    assert md5_hash == expected
+
+@pytest.mark.parametrize(["removeHs", "prerender", "expected"], [
+    (True, True, "acf7cf7cd5cfaa5eaf4a4f257e290e49"),
+    (False, True, ("2d30379732d312409d1d93ad5b3c3e60"
+                   if rdkit_version == "2020.03.1"
+                   else "b01c2146523b8a006bb34621839668eb")),
+    (True, False, "d8227a9948b6f801193085e5a8b257a2"),
+    (False, False, "e7b837b4d0c5f1ed8e1cd79734b6845e"),
+])
+def test_removeHs(driver, df, removeHs, prerender, expected):
+    useSVG = not prerender
+    mol = df["mol"][0]
+    mol.ClearProp("SMILES")
+    mols = [Chem.AddHs(mol)]
+    grid = mols2grid.MolGrid.from_mols(
+        mols, removeHs=removeHs, prerender=prerender, useSVG=useSVG, use_coords=False)
+    doc = get_doc(grid, {"n_rows": 1, "substruct_highlight": False})
+    driver.get(doc)
+    if not prerender:
+        driver.wait_for_img_load()
+    if useSVG:
+        md5_hash = driver.get_svg_md5_hash()
+    else:
+        md5_hash = driver.get_png_md5_hash()
+    assert md5_hash == expected
+
+@pytest.mark.parametrize(["kwargs", "expected"], [
+    (dict(addAtomIndices=True), "d9611dd45ff82ab73cf31dfc3a33706d"),
+    (dict(fixedBondLength=10), "1dfa2a5c43022ce66a2ff26a517698cf"),
+    (dict(atomColourPalette={6: (0, .8, .8)}), "91d114d421fc6a05e4537d19ad4b8324"),
+    (dict(legend="foo"), "bb5bc6ae713cfaea4e6374fc35e4be4e"),
+])
+def test_moldrawoptions(driver, df, kwargs, expected):
+    grid = get_grid(df, **kwargs)
+    doc = get_doc(grid, dict(n_rows=1, n_cols=1, subset=["img"]))
+    driver.get(doc)
+    driver.wait_for_img_load()
+    md5_hash = driver.get_svg_md5_hash()
+    assert md5_hash == expected
+
+def test_hover_color(driver, grid):
+    doc = get_doc(grid, {"hover_color": "red"})
+    driver.get(doc)
+    (ActionChains(driver)
+        .move_to_element(driver.find_by_css_selector("#mols2grid .cell"))
+        .perform())
+    color = driver.execute_script(f"return getComputedStyle(document.querySelector('#mols2grid .cell')).getPropertyValue('background-color');")
+    assert color == "rgb(255, 0, 0)"
+
+@flaky(max_runs=3, min_passes=1)
+def test_tooltip(driver, grid):
+    doc = get_doc(grid, {"tooltip": ["_Name"]})
+    driver.get(doc)
+    driver.wait_for_img_load()
+    (ActionChains(driver)
+        .move_to_element(driver.find_by_css_selector("#mols2grid .cell .data-img *"))
+        .perform())
+    tooltip = driver.find_by_css_selector('div.popover[role="tooltip"]')
+    el = tooltip.find_element_by_class_name("popover-body")
+    assert el.get_attribute("innerHTML") == "<strong>_Name</strong>: 3-methylpentane"
+
+@flaky(max_runs=3, min_passes=1)
+def test_tooltip_trigger(driver, grid):
+    doc = get_doc(grid, {"tooltip": ["_Name"], "tooltip_trigger": "click"})
+    driver.get(doc)
+    driver.wait_for_img_load()
+    (ActionChains(driver)
+        .move_to_element(driver.find_by_css_selector("#mols2grid .cell .data-img *"))
+        .perform())
+    with pytest.raises(NoSuchElementException):
+        driver.find_element_by_css_selector('div.popover[role="tooltip"]')
+    (ActionChains(driver)
+        .click(driver.find_by_css_selector("#mols2grid .cell .data-img *"))
+        .perform())
+    tooltip = driver.find_by_css_selector('div.popover[role="tooltip"]')
+    el = tooltip.find_element_by_class_name("popover-body")
+    assert el.get_attribute("innerHTML") == "<strong>_Name</strong>: 3-methylpentane"
+
+@flaky(max_runs=3, min_passes=1)
+def test_tooltip_fmt(driver, grid):
+    doc = get_doc(grid, {"tooltip": ["_Name"], "tooltip_fmt": "<em>{value}</em>"})
+    driver.get(doc)
+    driver.wait_for_img_load()
+    (ActionChains(driver)
+        .move_to_element(driver.find_by_css_selector("#mols2grid .cell .data-img *"))
+        .perform())
+    tooltip = driver.find_by_css_selector('div.popover[role="tooltip"]')
+    el = tooltip.find_element_by_class_name("popover-body")
+    assert el.get_attribute("innerHTML") == "<em>3-methylpentane</em>"
+
+def test_tooltip_not_in_subset(driver, grid):
+    doc = get_doc(grid, {"tooltip": ["_Name"], "subset": ["ID", "img"]})
+    driver.get(doc)
+    driver.wait_for_img_load()
+    (ActionChains(driver)
+        .move_to_element(driver.find_by_css_selector("#mols2grid .cell .data-img *"))
+        .perform())
+    tooltip = driver.find_by_css_selector('div.popover[role="tooltip"]')
+    el = tooltip.find_element_by_class_name("popover-body")
+    assert el.get_attribute("innerHTML") == "<strong>_Name</strong>: 3-methylpentane"
+
+@flaky(max_runs=3, min_passes=1)
+def test_style(driver, grid):
+    doc = get_doc(grid, {
+        "tooltip": ["_Name"],
+        "style": {
+            "__all__": lambda x: "color: red",
+            "_Name": lambda x: "color: blue",
+    }})
+    driver.get(doc)
+    driver.wait_for_img_load()
+    el = driver.find_by_css_selector("#mols2grid .cell")
+    assert el.value_of_css_property("color") == "rgb(255, 0, 0)"
+    el = driver.find_by_css_selector("#mols2grid .cell .data-_Name")
+    assert el.value_of_css_property("color") == "rgb(0, 0, 255)"
+    (ActionChains(driver)
+        .move_to_element(driver.find_by_css_selector("#mols2grid .cell .data-img *"))
+        .perform())
+    tooltip = driver.find_by_css_selector('div.popover[role="tooltip"]')
+    el = tooltip.find_element_by_class_name("popover-body")
+    assert el.get_attribute("innerHTML") == '<strong>_Name</strong>: <span style="color: blue">3-methylpentane</span>'
+
+@flaky(max_runs=3, min_passes=1)
+def test_transform(driver, grid):
+    doc = get_doc(grid, {
+        "tooltip": ["_Name"],
+        "transform": {"_Name": lambda x: x.upper()}
+    })
+    driver.get(doc)
+    name = driver.find_by_css_selector("#mols2grid .cell .data-_Name")
+    assert name.text == "3-METHYLPENTANE"
+    driver.wait_for_img_load()
+    (ActionChains(driver)
+        .move_to_element(driver.find_by_css_selector("#mols2grid .cell .data-img *"))
+        .pause(.5)
+        .perform())
+    tooltip = driver.find_by_css_selector('div.popover[role="tooltip"]')
+    el = tooltip.find_element_by_class_name("popover-body")
+    assert el.get_attribute("innerHTML") == '<strong>_Name</strong>: 3-METHYLPENTANE'
+
+@flaky(max_runs=3, min_passes=1)
+def test_transform_style_tooltip(driver, grid):
+    doc = get_doc(grid, {
+        "tooltip": ["_Name"],
+        "transform": {"_Name": lambda x: "foo"},
+        "style": {
+            "__all__": lambda x: "background-color: red",
+            "_Name": lambda x: "color: green" if x == "foo" else "color: blue",
+    }})
+    driver.get(doc)
+    driver.wait_for_img_load()
+    cell = driver.find_by_css_selector("#mols2grid .cell")
+    assert cell.value_of_css_property("background-color") == "rgb(255, 0, 0)"
+    name = cell.find_element_by_class_name("data-_Name")
+    assert name.text == "foo"   
+    (ActionChains(driver)
+        .move_to_element(driver.find_by_css_selector("#mols2grid .cell .data-img *"))
+        .pause(.5)
+        .perform())
+    tooltip = driver.find_by_css_selector('div.popover[role="tooltip"]')
+    el = tooltip.find_element_by_class_name("popover-body")
+    assert el.get_attribute("innerHTML") == '<strong>_Name</strong>: <span style="color: blue">foo</span>'
+
+@pytest.mark.parametrize("selection", [True, False])
+def test_callback_js(driver, grid, selection):
+    doc = get_doc(grid, {
+        "callback": "$('#mols2grid .cell .data-_Name').html('foo')",
+        "selection": selection
+    })
+    driver.get(doc)
+    driver.wait_for_img_load()
+    (ActionChains(driver)
+        .click(driver.find_by_css_selector("#mols2grid .cell .data-img"))
+        .pause(.5)
+        .perform())
+    el = driver.find_by_css_selector("#mols2grid .cell .data-_Name")
+    assert el.text == "foo"
+
+def test_sort_by(driver, grid):
+    doc = get_doc(grid, {"sort_by": "_Name"})
+    driver.get(doc)
+    el = driver.find_by_css_selector("#mols2grid .cell .data-_Name")
+    assert el.text == "1,1,2,2-tetrachloroethane"
+
+def test_sort_button(driver, html_doc):
+    driver.get(html_doc)
+    driver.wait_for_img_load()
+    (ActionChains(driver)
+        .click(driver.find_by_id("sortDropdown"))
+        .pause(.3)
+        .click(driver.find_by_css_selector('button.sort-btn[data-name="data-_Name"]'))
+        .perform())
+    el = driver.find_by_css_selector("#mols2grid .cell .data-_Name")
+    assert el.text == "1,1,2,2-tetrachloroethane"
+    (ActionChains(driver)
+        .click(driver.find_by_id("sortDropdown"))
+        .pause(.3)
+        .click(driver.find_by_css_selector('button.sort-btn[data-name="data-_Name"]'))
+        .perform())
+    el = driver.find_by_css_selector("#mols2grid .cell .data-_Name")
+    assert el.text == "tetrachloromethane"
+
+@pytest.mark.parametrize(["substruct_highlight", "expected"], [
+    (True,  "6a1ae3fc97cbe003c86db3c624e5ff09"),
+    (False, "4884c75977e8c3e0fa2e6a1607275406")
+])
+def test_substruct_highlight(driver, grid, substruct_highlight, expected):
+    doc = get_doc(grid, {"n_rows": 1, "substruct_highlight": substruct_highlight})
+    driver.get(doc)
+    text_box = driver.find_by_id("searchbar")
+    driver.wait_for_img_load()
+    (ActionChains(driver)
+        .click(driver.find_by_id("searchBtn"))
+        .pause(.5)
+        .click(driver.find_by_id("smartsSearch"))
+        .send_keys_to_element(text_box, "CC(I)C")
+        .key_up(Keys.TAB)
+        .perform())
+    md5_hash = driver.get_svg_md5_hash()
+    assert md5_hash == expected
+
+def test_substruct_clear_removes_highlight(driver, grid):
+    doc = get_doc(grid, {"n_rows": 1, "substruct_highlight": True})
+    driver.get(doc)
+    driver.wait_for_img_load()
+    (ActionChains(driver)
+        .click(driver.find_by_id("searchBtn"))
+        .pause(.5)
+        .click(driver.find_by_id("smartsSearch"))
+        .send_keys_to_element(driver.find_by_id("searchbar"), "C")
+        .key_up(Keys.TAB)
+        .perform())
+    md5_hash_hl = driver.get_svg_md5_hash()
+    driver.find_by_id("searchbar").clear()
+    (ActionChains(driver)
+        .send_keys_to_element(driver.find_by_id("searchbar"), Keys.BACKSPACE)
+        .key_up(Keys.TAB)
+        .perform())
+    md5_hash = driver.get_svg_md5_hash()
+    assert md5_hash != md5_hash_hl
+    assert md5_hash == "d8227a9948b6f801193085e5a8b257a2"
+
+def test_smarts_to_text_search_removes_highlight(driver, grid):
+    doc = get_doc(grid, {"n_rows": 1, "substruct_highlight": True})
+    driver.get(doc)
+    driver.wait_for_img_load()
+    (ActionChains(driver)
+        .click(driver.find_by_id("searchBtn"))
+        .pause(.5)
+        .click(driver.find_by_id("smartsSearch"))
+        .send_keys_to_element(driver.find_by_id("searchbar"), "I")
+        .key_up(Keys.TAB)
+        .perform())
+    md5_hash_hl = driver.get_svg_md5_hash()
+    (ActionChains(driver)
+        .click(driver.find_by_id("searchBtn"))
+        .pause(.5)
+        .click(driver.find_by_id("txtSearch"))
+        .send_keys_to_element(driver.find_by_id("searchbar"), "odopropane")
+        .key_up(Keys.TAB)
+        .perform())
+    md5_hash = driver.get_svg_md5_hash()
+    assert md5_hash != md5_hash_hl
+    assert md5_hash == "4884c75977e8c3e0fa2e6a1607275406"
+
+def test_filter(driver, grid):
+    doc = get_doc(grid, {})
+    driver.get(doc)
+    mask = grid.dataframe["_Name"].str.contains("iodopropane")
+    ids = grid.dataframe.loc[mask]["mols2grid-id"].to_list()
+    filter_code = env.get_template('js/filter.js').render(
+        grid_id = grid._grid_id,
+        ids = ids)
+    driver.wait_for_img_load()
+    driver.execute_script(filter_code)
+    el = driver.find_by_css_selector("#mols2grid .cell .data-_Name")
+    assert el.text == "2-iodopropane"
+
+def test_subset_gives_rows_order(driver, grid):
+    subset = ["_Name", "img", "ID"]
+    doc = get_doc(grid, {"subset": subset, "n_rows": 1})
+    driver.get(doc)
+    driver.wait_for_img_load()
+    cell = driver.find_by_css_selector("#mols2grid .cell")
+    elements = cell.find_elements_by_class_name("data")
+    for i, el in enumerate(elements):
+        class_list = el.get_attribute("class").split()
+        name = [x.replace("data-", "") for x in class_list
+                if x.startswith("data-")][0]
+        if i == len(elements) - 1:
+            assert name == "SMILES"
+            assert el.value_of_css_property("display") == "none"
+        else:
+            assert name == subset[i]
diff --git a/tests/test_molgrid.py b/tests/test_molgrid.py
index 46c08fa..a32d418 100644
--- a/tests/test_molgrid.py
+++ b/tests/test_molgrid.py
@@ -17,7 +17,7 @@ def _make_df():
 default_df = _make_df()
 
 def _make_grid(**kwargs):
-    kwargs["mol_col"] = kwargs.pop("mol_col", "mol")
+    kwargs.setdefault("mol_col", "mol")
     return MolGrid(default_df, **kwargs)
 
 @pytest.fixture(scope="module")
@@ -26,6 +26,10 @@ def df():
 
 @pytest.fixture(scope="module")
 def grid():
+    return _make_grid(prerender=True)
+
+@pytest.fixture(scope="module")
+def grid_otf():
     return _make_grid()
 
 @pytest.fixture(scope="module")
@@ -48,7 +52,7 @@ def test_uses_svg(grid):
     assert "<svg " in data[:100]
 
 def test_uses_png():
-    grid = _make_grid(useSVG=False)
+    grid = _make_grid(useSVG=False, prerender=True)
     assert grid.useSVG == False
     data = grid.dataframe.loc[0, "img"]
     assert data.startswith('<img src="data:image/png;base64')
@@ -85,7 +89,7 @@ def test_remove_hydrogens():
     mol =  Chem.MolFromSmiles("C")
     mol = Chem.AddHs(mol)
     data = [{"mol": mol, "ID": 0}]
-    grid = MolGrid(data, mol_col="mol", removeHs=True)
+    grid = MolGrid(data, mol_col="mol", removeHs=True, prerender=True)
     new = grid.dataframe.loc[0, "mol"]
     assert new.GetNumAtoms() == 1
 
@@ -94,31 +98,31 @@ def test_remove_coordinates():
     mol = Chem.AddHs(mol)
     Compute2DCoords(mol)
     data = [{"mol": mol, "ID": 0}]
-    grid = MolGrid(data, mol_col="mol", use_coords=False)
+    grid = MolGrid(data, mol_col="mol", use_coords=False, prerender=True)
     new = grid.dataframe.loc[0, "mol"]
     with pytest.raises(ValueError, match="Bad Conformer Id"):
         new.GetConformer()
 
 def test_generate_smi_from_mol_col(df):
     df = df.drop(columns=["SMILES"])
-    grid = MolGrid(df, mol_col="mol")
+    grid = MolGrid(df, mol_col="mol", prerender=True)
     assert "SMILES" in grid.dataframe.columns
 
 def test_index_when_none_molecules(df):
     df = df.copy()
     df.loc[1, "mol"] = None
-    grid = MolGrid(df, mol_col="mol")
+    grid = MolGrid(df, mol_col="mol", prerender=True)
     assert len(grid.dataframe) == len(df) - 1
     assert_equal(grid.dataframe["mols2grid-id"].values, [0, 2, 3, 4])
 
 def test_custom_moldrawoptions():
     opts = Draw.MolDrawOptions()
     opts.fixedBondLength = 42
-    grid = _make_grid(MolDrawOptions=opts)
+    grid = _make_grid(MolDrawOptions=opts, prerender=True)
     assert grid.MolDrawOptions.fixedBondLength == 42
 
 def test_moldrawoptions_as_kwargs():
-    grid = _make_grid(fixedBondLength=42)
+    grid = _make_grid(fixedBondLength=42, prerender=True)
     assert grid.MolDrawOptions.fixedBondLength == 42
 
 def test_update_current_grid_on_init():
@@ -137,12 +141,21 @@ def test_from_mols_custom_mol_col(mols):
     assert "mol" not in grid.dataframe.columns
 
 @pytest.mark.parametrize(["param", "attr", "value"], [
+    ("name", "_grid_id", "foobar"),
+    ("coordGen", "prefer_coordGen", False),
+    ("removeHs", "removeHs", True),
     ("useSVG", "useSVG", False),
+    ("use_coords", "use_coords", False),
     ("size", "img_size", (42, 42)),
-    ("name", "_grid_id", "foobar"),
+    ("prerender", "prerender", True),
+    ("smiles_col", "smiles_col", "smi"),
+    ("mol_col", "mol_col", "rdmol"),
 ])
 def test_from_mols_kwargs(mols, param, attr, value):
-    grid = MolGrid.from_mols(mols, **{param: value})
+    if param in ["useSVG"]:
+        grid = MolGrid.from_mols(mols, prerender=True, **{param: value})
+    else:
+        grid = MolGrid.from_mols(mols, **{param: value})
     assert getattr(grid, attr) == value
 
 def test_from_sdf():
@@ -150,17 +163,26 @@ def test_from_sdf():
     assert "mol" in grid.dataframe.columns
 
 def test_from_sdf_custom_mol_col():
-    grid = MolGrid.from_sdf(sdf, mol_col="rdmol")
+    grid = MolGrid.from_sdf(sdf, mol_col="rdmol", prerender=True)
     assert "rdmol" in grid.dataframe.columns
     assert "mol" not in grid.dataframe.columns
 
 @pytest.mark.parametrize(["param", "attr", "value"], [
+    ("name", "_grid_id", "foobar"),
+    ("coordGen", "prefer_coordGen", False),
+    ("removeHs", "removeHs", True),
     ("useSVG", "useSVG", False),
+    ("use_coords", "use_coords", False),
     ("size", "img_size", (42, 42)),
-    ("name", "_grid_id", "foobar"),
+    ("prerender", "prerender", True),
+    ("smiles_col", "smiles_col", "smi"),
+    ("mol_col", "mol_col", "rdmol"),
 ])
 def test_from_sdf_kwargs(param, attr, value):
-    grid = MolGrid.from_sdf(sdf, **{param: value})
+    if param in ["useSVG"]:
+        grid = MolGrid.from_sdf(sdf, prerender=True, **{param: value})
+    else:
+        grid = MolGrid.from_sdf(sdf, **{param: value})
     assert getattr(grid, attr) == value
 
 def test_template(grid):
@@ -174,19 +196,19 @@ def test_template_setter(grid):
 def test_substruct_highlight():
     mol = Chem.MolFromSmiles("C")
     mol.__sssAtoms = [0]
-    grid = MolGrid.from_mols([mol])
+    grid = MolGrid.from_mols([mol], prerender=True)
     svg = grid.dataframe.loc[0, "img"]
     assert "<ellipse" in svg
 
 def test_mol_to_img_svg():
     mol = Chem.MolFromSmiles("C")
-    grid = MolGrid.from_mols([mol])
+    grid = MolGrid.from_mols([mol], prerender=True)
     img = grid.mol_to_img(mol)
     assert "<svg " in img[:100]
 
 def test_mol_to_img_png():
     mol = Chem.MolFromSmiles("C")
-    grid = MolGrid.from_mols([mol])
+    grid = MolGrid.from_mols([mol], prerender=True)
     grid.useSVG = False
     grid._MolDraw2D = Draw.MolDraw2DCairo
     img = grid.mol_to_img(mol)
@@ -199,22 +221,22 @@ def test_get_selection(df):
     assert_equal(new.values,
                  df.iloc[0:1].values)
 
-def test_save(grid):
+def test_save(grid_otf):
     with NamedTemporaryFile("w", suffix=".html") as f:
-        grid.save(f.name)
+        grid_otf.save(f.name)
         assert Path(f.name).is_file()
 
 @pytest.mark.parametrize("kind", ["pages", "table"])
-def test_render(grid, kind):
-    grid.render(template=kind)
+def test_render(grid_otf, kind):
+    grid_otf.render(template=kind)
 
-def test_render_wrong_template(grid):
+def test_render_wrong_template(grid_otf):
     with pytest.raises(ValueError, match="template='foo' not supported"):
-        grid.render(template="foo")
+        grid_otf.render(template="foo")
 
 @pytest.mark.parametrize("kwargs", [
     dict(),
-    dict(subset=["ID"]),
+    dict(subset=["img"]),
     dict(tooltip=["ID"]),
     dict(selection=False),
     dict(style={"ID": lambda x: "color: red" if x == 1 else ""}),
@@ -222,9 +244,10 @@ def test_render_wrong_template(grid):
     dict(custom_css="* {color: red;}"),
     dict(callback="console.log(JSON.stringify(data));"),
     dict(custom_header='<script src="https://unpkg.com/@rdkit/rdkit@2021.3.2/Code/MinimalLib/dist/RDKit_minimal.js"></script>'),
+    dict(substruct_highlight=False),
 ])
-def test_integration_pages(grid, kwargs):
-    grid.to_pages(**kwargs)
+def test_integration_pages(grid_otf, kwargs):
+    grid_otf.to_pages(**kwargs)
 
 @pytest.mark.parametrize("kwargs", [
     dict(),
@@ -237,15 +260,57 @@ def test_integration_pages(grid, kwargs):
 def test_integration_table(grid, kwargs):
     grid.to_table(**kwargs)
 
-def test_python_callback(grid):
+def test_python_callback(grid_otf):
     def myfunc():
         pass
-    html = grid.to_pages(subset=["ID"], callback=myfunc)
+    html = grid_otf.to_pages(subset=["img"], callback=myfunc)
     assert "// call custom python callback" in html
     assert "// no kernel detected for callback" in html
 
-def test_python_callback_lambda(grid):
+def test_python_callback_lambda(grid_otf):
     with pytest.raises(TypeError, match="Lambda functions are not supported"):
-        grid.to_pages(subset=["ID"], callback=lambda x: None)
+        grid_otf.to_pages(subset=["img"], callback=lambda x: None)
+
+def test_cache_selection():
+    grid = _make_grid(name="cache")
+    register.add_selection("cache", [0], ["CCO"])
+    grid = _make_grid(name="cache", cache_selection=True)
+    assert hasattr(grid, "_cached_selection")
+    assert register.get_selection("cache") == grid._cached_selection
+
+def test_cache_no_init():
+    grid = _make_grid(name="new_cache", cache_selection=True)
+    assert hasattr(grid, "_cached_selection")
+    assert grid._cached_selection is None
+    assert "new_cache" in register.list_grids()
+
+def test_no_cache_selection():
+    grid = _make_grid(name="no_cache", cache_selection=False)
+    assert grid._cached_selection is None
+    assert "no_cache" in register.list_grids()
+
+def test_onthefly_render_png_error(df):
+    with pytest.raises(ValueError,
+                       match="On-the-fly rendering of PNG images not supported"):
+        MolGrid(df, prerender=False, useSVG=False)
+
+def test_substruct_highlight_prerender_error(grid):
+    with pytest.raises(ValueError,
+                       match="Cannot highlight substructure search with "
+                             "prerendered images"):
+        grid.display()
+
+def test_use_coords_onthefly_error():
+    with pytest.raises(ValueError,
+                       match="Cannot use coordinates with on-the-fly rendering"):
+        _make_grid(use_coords=True, prerender=False)
+
+def test_sort_by_not_in_subset_or_tooltip(grid_otf):
+    with pytest.raises(ValueError,
+                       match="'_Name' is not an available field"):
+        grid_otf.to_pages(subset=["ID", "img"], sort_by="_Name")
 
-# TODO: test filters and display
\ No newline at end of file
+def test_subset_without_img_error(grid_otf):
+    with pytest.raises(KeyError,
+                       match="Please add the 'img' field in the `subset` parameter"):
+        grid_otf.display(subset=["_Name"])
\ No newline at end of file
diff --git a/tests/test_select.py b/tests/test_select.py
index 6d9e126..a3f8e43 100644
--- a/tests/test_select.py
+++ b/tests/test_select.py
@@ -1,3 +1,4 @@
+from threading import Thread
 import pytest
 import mols2grid as mg
 from mols2grid.select import register
@@ -47,14 +48,15 @@ def test_unset_selection():
 
 def test_add_selection():
     register._init_grid("bar")
-    register.add_selection("foo", 0, "C")
+    register.add_selection("foo", [0, 1], ["C", "CC"])
     assert register.get_selection()[0] == "C"
+    assert register.get_selection()[1] == "CC"
 
 def test_del_selection():
-    register._set_selection(0, "C")
     register._init_grid("bar")
-    register.del_selection("foo", 0)
-    assert register.get_selection() == {}
+    register.add_selection("foo", [0, 1], ["C", "CC"])
+    register.del_selection("foo", [1])
+    assert register.get_selection() == {0: "C"}
 
 def test_get_selection():
     assert register.get_selection() == {}
@@ -74,3 +76,18 @@ def test_warn_selection():
     with pytest.warns(UserWarning,
                       match="`mols2grid.selection` is deprecated"):
         len(mg.selection)
+
+def test_lock(n=1e6):
+    N = int(n)
+    x = list(range(N))
+    def func(name):
+        register.add_selection(name, x, x)
+    register._init_grid("bar")
+    t1 = Thread(target=func, args=("foo",))
+    t2 = Thread(target=func, args=("bar",))
+    t1.start()
+    t2.start()
+    t1.join()
+    t2.join()
+    assert len(register.get_selection("foo")) == N
+    assert len(register.get_selection("bar")) == N
diff --git a/tests/test_utils.py b/tests/test_utils.py
index cdf3b6a..4f67831 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -1,3 +1,5 @@
+from tempfile import NamedTemporaryFile
+import gzip
 import pytest
 import pandas as pd
 from rdkit import RDConfig, Chem
@@ -23,7 +25,7 @@ def func():
     (["Activity"], "{value}", {}, {"Activity": lambda x: f"{x:.2f}"},
      "42.01"),
     (["Activity"], "{key}: {value}",
-     {"Activity": lambda x: "color: red" if x.endswith("1") else ""},
+     {"Activity": lambda x: "color: red" if x > 40 else ""},
      {"Activity": lambda x: f"{x:.2f}"},
      'Activity: <span style="color: red">42.01</span>'),
 ])
@@ -64,9 +66,8 @@ def test_mol_to_record():
         else:
             mol.SetProp(prop, value)
     new = utils.mol_to_record(mol)
-    for default_col in ["mol", "SMILES"]:
-        assert default_col in new.keys()
-        new.pop(default_col)
+    assert "mol" in new.keys()
+    new.pop("mol")
     assert new == props
 
 def test_mol_to_record_none():
@@ -109,6 +110,15 @@ def test_sdf_to_dataframe_custom_mol_col():
     assert "mol" not in df.columns
     assert "foo" in df.columns
 
+def test_sdf_to_df_gz():
+    with NamedTemporaryFile("wb", suffix=".gz") as tf, open(sdf, "rb") as fi:
+        gz = gzip.compress(fi.read(), compresslevel=1)
+        tf.write(gz)
+        tf.flush()
+        df = utils.sdf_to_dataframe(tf.name).drop(columns=["mol"])
+        ref = utils.sdf_to_dataframe(sdf).drop(columns=["mol"])
+        assert (df == ref).values.all()
+
 def test_remove_coordinates():
     mol = Chem.MolFromSmiles("CCO")
     Compute2DCoords(mol)