marinegor.dev

marimo: Python notebooks that are just Python

Thu, 24 Sep 2026 00:00:00 +0000

marimo is a Python notebook, like Jupyter, but built around a few decisions that fix most of what annoys me about notebooks.

Why it’s cool

Reactive: marimo knows which cells use which variables. Change a cell or move a slider, and every cell that depends on it re-runs; delete a cell, and its variables are gone. No hidden state, no “restart and run all” to check that it actually works top to bottom.
Just a .py file: a notebook is a plain Python script. It diffs nicely in git, can be linted and formatted like any other code (this repo runs ruff and ty on it), imported as a module, or run with python notebook.py.
Dependencies inside the notebook: packages are declared inline (PEP 723), so marimo edit --sandbox notebook.py builds an isolated environment with uv and you just run it.
Apps for free: UI elements (sliders, dropdowns, forms, …) are regular Python values, so any notebook can be shown as an interactive app with marimo run.
Runs in the browser: via WebAssembly and Pyodide, with no server at all. That’s what makes the embed below possible on a static site like this one.

A live example

Below is the cellular automaton art notebook from the marimo gallery. It’s the same one used in the marimo embedding docs. It runs entirely in your browser, so the first load takes a few seconds while Python itself is downloaded. Then click Make art!:

This post embeds an interactive marimo notebook (cellular-automaton-art), which doesn't run in feed readers: open the post to try it, or see its source.

The notebook lives in this site’s repo as a plain .py file, and is also served as-is. The marimo shortcode reads it at build time and passes its code to marimo.app in the iframe URL. To play with it locally:

uvx marimo edit --sandbox https://marinegor.dev/notebooks/cellular-automaton-art.py

Split-and-combine parallelization for MDAnalysis AnalysisBase subclasses

Fri, 16 Aug 2024 00:00:00 +0000

🔗 MDAnalysis PR #4162

You can find all about this project by gsoc tag, but basically, you can now do ArbitraryAnalysisSubclass.run(backend='multiprocessing'), and it’ll run automatically in multiple processes.

You can read more about parallelization on docs, especially some limitations of the current approach.

Secondary structure analysis with DSSP for MDAnalysis

Tue, 18 Jun 2024 00:00:00 +0000

🔗 MDAnalysis PR #4304

A fun project emerged during MDAnalysis 2023 UGM hackathon. Basically, it adds a built-in secondary structure analysis, with no dependencies. The algorithm is based on Kabsch et al, and implementation is borrowed from pyDSSP. I slightly modified the implementation, making it depend only on numpy (previously it depended on einops package).

For details, see documentation.

A subgroup of light-driven sodium pumps with an additional Schiff base counterion

Wed, 10 Apr 2024 00:00:00 +0000

📄 Nature Communications (2024)

A project led by super talented Kirill Kovalev. Whereas it’s always a pleasure to work with Kirill, it’s also a great study that combines multiple techniques, both functional and structural, to study a novel rhodopsin. For me personally, it’s a first paper that confirms I can do cryoEM data processing🥶

On the methods of studying protein-ligand interaction dynamics

Mon, 12 Feb 2024 00:00:00 +0000

📄 PhD thesis, University of Groningen (2024)

My thesis, written under supervision of Albert Guskov, Valentin Borshchevskiy and Dirk Slotboom. The text is open access, and is basically composed of few articles mentioned below, plus some extra nice acknowledgements in multiple languages.

Regression-based Accelerated Virtual Screening

Fri, 29 Dec 2023 00:00:00 +0000

Article link

When doing large-scale virtual screening, you have a slow oracle that tells you the true value of each molecule (in our case, it’s small molecule docking). Point is, the oracle is slow – in our case, around 1 CPU-sec per molecule (depends heavily on pocket size and ligand sampling depth), which is a lot if you have $10^9$ molecules.

Main idea here, described also before us, is to do an active learning, using the docking algorithm as oracle. Then we can dock a small batch, then train a model to predict the docking scores, and then dock only those that have good scores. Potentially, we can also do that many times. It’s quite intriguing though that, despite the whole “accelerated virtual screening” started back in 2020, all the papers that I’ve seen use some heavy-weight model as a base one, i.e. the model that determines which molecules go to the next docking step.

Here we showed that you can actually use a linear regression (literally, sklearn.linear_model.LinearRegression) model, trained on Morgan fingerprints, instead, and also few slightly more clever tricks to make your retrieval rate better. Namely:

use smaller batch size – around 10_000 is perfect
re-use models – aggregate few (we did 20) last batches via MeanRank, i.e. take the mean across last 20 models predicted rank of each molecule. You can also do some exponential decay of that, if you want.
use relatively big fingerprints – this preprint by Martin et al benchmarks it much better
models like sklearn.svm.SVR could also be surprisingly good, though less stable across datasets

Our final active learinng model was following:

Morgan fingerprints with 1024 bits, radius 2, as features
sklearn.linear_model.LinearRegression as base learner
batch size 10_000
aggregation of last 20 models via mean ranking

With that, we could retrieve 70% of the top-0.05% of famous AmpC dataset from Lyu et al after screening only 2% of the library, whereas the message-passing neural networks approach retrieves around 80%. Kinda cool, given that we used a CPU-only workstation with only a modest SSD for our work.

Also, cool thing is that our approach allows us to pre-compute the fingerprints and store them somewhere, whereas MPNN-based one needs to be trained again and again for each new dataset, which is also kinda a disadvantage.

Regression-Based Active Learning for Accessible Acceleration of Ultra-Large Library Docking

Fri, 29 Dec 2023 00:00:00 +0000

📄 Journal of Chemical Information and Modeling (2023)

Amazing (conflict of interest: I am the last author) paper, where we basically show that deep learning isn’t needed for accelerated virtual screening.

For more info, please go to the article description

Custom Design of a Humidifier Chamber for In Meso Crystallization

Tue, 12 Dec 2023 00:00:00 +0000

📄 Crystal Growth & Design (2023)

A neat paper that shares an open-source design (as in “to assemble on your own if you have decent skills”) of a humidifier chamber for a Gryphon crystallization robot, that unfortunately doesn’t have such chamber.

MDAnalysis UGM 2023: MDAnalysis Parallelization @ GSOC

Wed, 27 Sep 2023 00:00:00 +0000

The content below is one-to-one a content of my MDAnalysis User Meeting in September 2023.

Parallelization of MDAnalysis

Egor Marin, Google Summer of Code contributor

marinegor@fastmail.com

Motivation

T(insight) = T(simulation) + T(analysis)
pmda: 30⭐️, 21🍴 (forks)
all MDAnalysis instances are parallelizable (courtesy @yuxuanzhuang)

Goals

Change AnalysisBase in a way that…

Introduces multiple job support
- local – using multiprocessing?
- non-local – using dask?
Does not break existing code (subclasses, dependencies, etc)
Allows customization
- add your own execution backend
- write your own parallelizable classes

Sneak peak: results

>>> from MDAnalysis.analysis.rms import RMSD
>>> u, ref = ...

>>> % time long_run = RMSD(u, ref)
Wall time: 4 min

>>> % time long_run = RMSD(u, ref, backend='multiprocessing', n_workers=8)
Wall time: 50 sec

Sneak peak: results

>>> from MDAnalysis.analysis.rdf import InterRDF
>>> s1, s2 = ...

>>> % time long_run = InterRDF(s1, s2)
Wall time: 1 h

>>> % time long_run = InterRDF(s1, s2, backend='multiprocessing', n_workers=8)
Wall time: 8 min

Implementation

General AnalysisBase.run() protocol:

def run(self, 
        start=None, 
        stop=None, 
        step=None, 
        frames=None, ...):    
  
  # prepare frames -- check boundaries, setup reader
  self._setup_frames(self._trajectory, start=start, stop=stop, step=step, frames=frames)
    
  # initialize attributes with intermediate results
  self._prepare()

  # go through trajectory 
  for i, ts in enumerate(self._sliced_trajectory, ...):
    self._single_frame()

  # convert intermediate results into final ones
  self._conclude()

Implementation

split-apply-combine:

split all frames into groups
process each group independently
combine results

Implementation

Additional methods of AnalysisBase:

“split”: _setup_computation_groups()
“apply”: _compute
“combine”: _get_aggregator & ResultsGroup.merge()
housekeeping: _configure_backend

Implementation

Additional classes: ResultsGroup

class ResultsGroup:
    def __init__(self, lookup: dict[str, Callable]): ...
    def merge(
        self, objects: Sequence[Results], require_all_aggregators: bool = True
    ) -> Results: ...

Implementation

Additional classes: ResultsGroup

>> from MDAnalysis.analysis.parallel import ResultsGroup, Results
>> group = ResultsGroup(lookup={'mass': lambda values: np.mean(values)}
>> obj1 = Results(mass=1)
>> obj2 = Results(mass=3)
>> group.merge([obj1, obj2])
{'mass': 2.0}

Implementation

Additional classes: ResultsGroup

# skip attributes you don't want to aggregate
>>> lookup = {'mass': ResultsGroup.float_mean, 'trajectory': Trajectory}
>>> group = ResultsGroup(lookup)
>>> objects = [Results(mass=1, trajectory=None), Results(mass=3, trajectory=Trajectory(...))]
>>> group.merge(objects, require_all_aggregators=False)
{'mass': 2.0}

Implementation

Additional classes: BackendBase

from MDAnalysis.parallel import BackendBase


class CustomBackend(BackendBase):
    def __init__(self, some_resource):
        self.some_resource = some_resource

    def apply(self, func: Callable[T, R], computations: list[T]) -> list[R]:
        results = [self.some_resource.do_compute(func, task) for task in computations]
        return results

Implementation

Additional classes: BackendBase

>>> backend_instance = CustomBackend(some_resource=...)
>>> my_run = RMSD(u, ref)
>>> my_run.run(backend=backend_instance, unsafe=True)

Implementation

Conditions for the backend=...

Subclass.is_parallelizable() == True
backend: str|BackendBase
- if str, one of built-in backends (multiprocessing/dask)
- if BackendBase subclass, explicitly say unsafe=True

Add to your subclass

Example: RMSD

from MDAnalysis.analysis.parallel import ResultsGroup


class RMSD(AnalysisBase):
    @classmethod
    @property
    def available_backends(cls):
        return (
            "serial",
            "multiprocessing",
            "dask",
        )

    @classmethod
    @property
    def is_parallelizable(self):
        return True

    def _get_aggregator(self):
        return ResultsGroup(lookup={"rmsd": ResultsGroup.ndarray_vstack})

How fast

full-atom xtc (“lysozyme in water”)
two different drives, hdparm:
- HDD: 180 MB/sec
- SSD: 460 MB/sec
i9-9900K CPU @ 3.60GHz, 8 cores/socket

How fast

When should I use it

Definitive answer: benchmark it yourself!

%%time in jupyter notebook
step=MANY or start=0, stop=10_000 to achieve ~1-2 minute runtime with serial scheduler

When should I use it

YourClass.is_parallelizable() for your class
_single_frame() slower than reading from disk

Future features

🔜 add to all subclasses
- have separate issues for complicated ones?
🔜 add tutorials
- how to use parallel backends
- how to implement your own backend subclasses

Future features

🤔 add distributed support
- rebrand pmda into MDAKit
- focus on parallel filesystems, etc
🤔 add shareable memory trajectories
- pre-load trajectory & analyze it
- efficient for practically all subclasses
🤔 add AnalysisCollection (ht @PicoCentauri)
- one frame read exactly once

Acknowledgements

Google Summer of Code 2023 program

Acknowledgements

Mentors:
- Yuxuan Zhuang @yuxuanzhuang
- Oliver Beckstein @orbeckst
- Rocco Meli @RMeli
💫 Paul Smith @p-j-smith
all contributors🫶

Taskfile + uv for running python tests

Fri, 15 Sep 2023 00:00:00 +0000

Glossary

What is a Taskfile

Like Makefile, but tasks are written in bash: link. Basically, it allows you to define custom, arbitrarily complex tasks, and then run then with task <taskname>.

What is `uv`

A super fast python package manager ~~that everyone’s been talking about~~. Just go check it out.

For me the primary usage for uv is replacing pip install ... with uv pip install ... – just makes things run faster (and with more clear error messages). However, here I’ll talk about a new feature called uv run – it allows you to run an arbitrary command without leaving current environment.

Taskfile example: MDAnalysis

Basically, the Taskfile:

version: '3'

tasks:
  clean:
    cmds:
      - rm -rf .venv
  ipython:
    cmds:
      - ipython
  build:
    cmds:
      - echo 'running build...'
      - uv venv -p 3.12
  rebuild:
    cmds:
      - task: clean
      - task: build
  install:
    cmds:
      - task: build
      - uv pip install -e './package[extra_formats,analysis,doc,parallel]' 
      - uv pip install -e './testsuite' pytest-xdist pytest-testmon pytest-random-order
      - uv pip install ipython ipykernel
  test:
    cmds:
      - task: build
      - task: install
      - uv run python -m pytest --random-order -n $(nproc --ignore=4) {{.CLI_ARGS}}
  check:
    cmds:
      - uv run python -m pytest --random-order --testmon -n $(nproc --ignore=4) {{.CLI_ARGS}}
  all:
    cmds:
      - task: clean
      - task: build
      - task: install
      - task: test

It allows you to build the full environment quickly (~2 minutes with cold cache), install the things I always install for tests but keep forgetting when re-installing the package (pytest-xdist and pytest-random-order), and run either all tests (task test) or only those that could’ve changed correctness since the last time I ran pytest (task check). The last one is useful when developing code.

Also, I can run ipython console, making sure that I’m running inside current environment, with task ipython. Pretty neat!

Don’t forget to add this Taskfile to your .git/info/exclude file in order not to commit it by accident.

GSOC-2023: summary

Wed, 23 Aug 2023 00:00:00 +0000

As you might know from my previous posts, during the summer of 2023 I’ve been working on MDAnalysis’s project during Google Summer of Code. Here I’ll summarize what I’ve done, how others can use it, and what changes will follow that in the MDAnalysis codebase in the near future.

A short description of the goals of the project.

One sentence: introduce parallel execution of analysis runs in MDAnalysis library. Somewhat good introduction I also gave here when writing a proposal for the project.

In more technical details, MDAnalysis library (as of v2.6.0) contains around 30 different subclasses that can perform various molecular trajectory analysis tasks, like calculating RMSD, RMSF, various contacts, density analysis, and more advanced tasks. 24 of these subclasses are children of AnalysisBase class. This base class is written in a way that allows subclass authors care about implementing only 3 methods:

_prepare() – how to initialize attributes for the analysis
_single_frame() – how to do analysis of a single molecular trajectory frame
_conclude() – how to transform intermediate results into final ones

With only these 3 methods implemented, by inheritance subclasses get run() method that will take care of reading the trajectory, storing the results, logging, etc.

The main goal was to re-write AnalysisBase.run() so that it can run in parallel on multiple processes, but make these changes invisible from the subclasses, i.e. not re-write any of their code.

What you did.

Altogether, the AnalysisBase.run() has changed by addition of the following methods:

_setup_computation_groups(): split frames into multiple parts for separate analysis
_compute(): run _single_frame on a list of frames, but without running _conclude
_get_aggregator(): get an object to aggregate the run results with, making them compatible with subsequent _conclude
class property available_backends(): get list of str values that describe available backends for a given subclass

I’ve also added ParallelExecutor and ResultsGroup classes that abstract away parallel execution and results aggregation, respectively. And finally, I added multiprocessing and dask/dask.distributed backends that reportedly speed up the analysis!

The current state.

Currently, changes to the AnalysisBase are almost finalized. One thing that holds it back is some CI/CD issues causing tests to timeout, but all AnalysisBase-related tests run both locally and on CI/CD system.

What’s left to do.

Optional thing suggested within my proposal was to actually add parallelization to the subclasses, and update tests accordingly. It turned out to be a tedious task, but finally the mechanism is finalized and described in MDAnalysisTests/analysis/conftest.py. It automatically generates each subclass fixtures for testing, and updating tests accordingly is fairly simple yet tedious. After updating all the tests, the library will be fully equipped with the parallel execution mechanisms for those classes that allow it.

What code got merged (or not) upstream.

The main changes are summarized in the main pull-request of the project. They mostly involve changes to package/analysis/base.py, as well as installation and CI/CD configuration files. Also, there are example changes in package/analysis/rms.py introducing parallelization into RMSD and RMSF subclasses, showcasing the changes to be made in order to add parallelization to a certain class.

How others can use it

Let’s imagine we’ve just learned that MDAnalysis now supports parallelization, and want to calculate RMSD of our large trajectory faster using multiple cores on our machine. Imagine we have a 16-core CPU Linux workstation with an SSD drive and a 1 us trajectory of lysozyme in water. Like this:

import MDAnalysis as mda
from MDAnalysis.analysis import rms

prefix = "./large_data/"
traj, top = f"{prefix}/md_0_1.xtc", f"{prefix}/md_0_1.gro"

u = mda.Universe(top, traj)

First we want to get a reference structure by taking the average of all frames. Like this:

from MDAnalysis.analysis.align import AverageStructure

avg = AverageStructure(mobile=u).run(backend="multiprocessing", n_workers=16)

but we get this:

ValueError                                Traceback (most recent call last)
Cell In[11], line 2
      1 from MDAnalysis.analysis.align import AverageStructure
----> 2 avg = AverageStructure(mobile=u).run(backend='multiprocessing', n_workers=16)
...
ValueError: backend=multiprocessing is not in self.available_backends=('local',) for class AverageStructure

which basically says we can use only backend='local' for the AverageStructure. Ok, let’s do that, but with a large step to save time:

avg = AverageStructure(mobile=u).run(step=100)
ref = avg.results.universe

and start our analysis run – RMSD for multiple selections to later compare them between each other:

groupselections = ("protein", "backbone", "name CA")

R = rms.RMSD(
    u,  # universe to align
    ref,  # reference universe or atomgroup
    groupselections=groupselections,
    # select="backbone",  # group to superimpose and calculate RMSD
)

If we start it with R.run(), we won’t even know when the run would finish. Luckily, we can add some verbosity with R.run(verbose=True) and see a nice tqdm progressbar that shows us that ETA of the whole analysis is around 4 minutes:

>>> R.run(verbose=True)
5%|▌         | 5062/100001 [00:13<04:15, 371.20it/s]

let’s try to speed it up now. Which backends do we have available?

>>> rms.RMSD.available_backends
('local', 'multiprocessing', 'dask', 'dask.distributed')

let’s try a built-in multiprocessing first:

>>> R.run(backend='multiprocessing', n_workers=4)
# CPU times: user 153 ms, sys: 74.2 ms, total: 227 ms
# Wall time: 1min 14s

ok, this is roughly 4 times faster! Amazing, roughly as we expected. Spoiler though – if we do it with 16 workers, we’ll see the total time around 40 seconds, so improvement saturates at some point.

But, we’ve lost something valuable when switching to multiprocessing – we don’t have a decent progressbar anymore. Luckily, we can use an amazing dask dashboard that allows us to monitor all tasks given to a particular dask.distributed cluster!

Let’s set up a cluster first:

from dask.distributed import Client, LocalCluster

cluster = LocalCluster(n_workers=8, threads_per_worker=1, memory_limit="30Gb")
client = Client(cluster)

and open the dashboard in our browser:

>>> cluster.dashboard_link
'http://127.0.0.1:8787/status'

Now, we’re ready to pass the pre-configured client as an argument to our R.run():

R.run(client=client)

unfortunately, we won’t see much progress – we can see that all tasks got spawned, but their status will change only upon completion, and we won’t get any intermediate progress report.

But luckily, there is a way to control that: in R.run() function, you can split the workload into an arbitrary number of parts with n_parts = ..., and upon completion of each of them dask would report that. Let’s do this:

R.run(client=client, n_parts=96)

now we’ll see intermediate progress as well as soon as each part gets completed, which is super helpful when trying to estimate the completion time.

Conclusion

We’ve now went through the essence of the MDAnalysis parallelization project, and learned how to use it in your analysis either by simply adding backend='multiprocessing', n_workers=..., or setting up your own dask cluster and submitting your jobs there.

Hopefully, this project will grow further and include all existing subclasses, as well as improving the speed (which, as we saw, saturates) and memory efficiency.

If you want to contribute to the project, stay tuned for the new issues on MDAnalysis github – there definitely will be some parallelization-related things for the future described there!

(again bi)weekly GSOC-4: finally about results aggretation

Mon, 14 Aug 2023 00:00:00 +0000

In the last blogpost, I explained how I decided to lift the parallelization class up, and how AnalysisBase.run() method changed after introducing this.

Here, I will (finally) talk about aggregating results from different worker objects, and how to make the implementation more explicit while not making people who create subclasses write a lot of boilerplate code.

What are we trying to do?

We’re not looking on how to aggregate results from independent AnalysisBase._compute() methods. Also, we want conclude() methods of subclasses to run as they used to, so after we run our aggregation, results and all other attributes should look exactly as if we’d run run(backend='local').

Looking at the _compute() method signature, we can see that it returns self (AnalysisBase). Also, the main object we want to return is whatever the type of results is. Usually it is MDAnalysis.analysis.Results object (basically, a fancy dict), although for some old classes it might be np.ndarray or just a list.

So, we’re looking for something like this:

def aggregate(remote_objects: list[AnalysisBase]) -> Results: ...

Note that we have an ordered sequence of remote results, and can rely on this fact when aggregating the result without the need to match frame indices with respective results.

First ideas

First thought that came at least to me looked like this: let’s take first object of remote_objects as a template, and based on how it looks like (=its type), try to do aggregation for it. Something like this:

def merge(remote_objects: list[AnalysisBase]) -> Results:
    template: Results | list | np.ndarray = remote_objects[0].results
    if isinstance(template, list):
        remote_results = [obj.results for obj in remote_objects]
        flat_results = flatten_arrays(objects_to_flatten)
    elif isinstance(template, np.ndarray):
        remote_results = np.array([obj.results for obj in remote_objects])
        flat_results = np.hstack(remote_results)
    elif isinstance(template, Results):
        ...
    else:
        raise ValueError("unknown results type")


def flatten_arrays(arrs: list[list]) -> list:
    return [obj for sublist in arrs for obj in sublist]

A little bit more complicated if isinstance(remote_results, Results). Since Results is basically a dictionary, we have to go over its keys and do proper aggregation for each of them:

... 
elif isinstance(template, Results):
    for key, obj_of_type in template.items():
        results_of_key = [obj.results[key] for obj in remote_objects]
        if isinstance(obj_of_type, list):
            flat_results = flatten_arrays(results_of_key)
        elif isinstance(obj_of_type, np.ndarray):
            flat_results = np.hstack(np.array(results_of_key))
        elif isinstance(obj_of_type, float):
            ...
        else:
            raise ValueError("couldn't find aggregation function")

Why first ideas are bad

The code above is probably already ugly for an experienced eye – for instance, we’re repeating ourselves quite a few times and have many nested if-elif-else parts that look super ugly. However, there’s a more important flaw here from the library point – if a subclass wants to have a custom aggregation function for its attribute, there is now way to conveniently reuse the base class method. Also, sometimes we can have different aggregation functions even if the objects have the same type – some ndarray-s have to be averaged, some simply summed, etc. The only other thing we have to distinguish results from each other is the attribute name, so we must base our aggregation on them.

To sum up, we actually don’t want to do any type matching (and ugly if-elif-else branches), but instead want to:

rely on the attribute name, not type
provide a reusable library of basic aggregation functions
allow users to easily create their own aggregation functions without copying our boilerplate

Good ideas: another new class

So we want to aggregate results based on attribute name, and also store somewhere a list of pre-determined functions that do this. We probably would be fine with a complicated function that stores everything, but since we want to also store different aggregation functions somewhere, let’s create a simple class which staticmethods would be different aggregation functions, and one useful method would be our merge from above:

class ResultsGroup:
    def merge(self, remote_objects: list[AnalysisBase]) -> Results: ...

    @staticmethod
    def flatten_arrays(arrs: list[list]) -> list:
        return [obj for sublist in arrs for obj in sublist]

    @staticmethod
    def ndarray_sum(arrs: list[np.ndarray]) -> np.ndarray:
        return np.array(arrs).sum(axis=0)

    @staticmethod
    def ndarray_stack(arrs: list[np.ndarray]) -> np.ndarray:
        return np.hstack(arrs)

    @staticmethod
    def ndarray_mean(arrs: list[np.ndarray]) -> np.ndarray:
        return np.array(arrs).mean(axis=0)

    ...

But how do we initialize the class? Well, we want to be able to call our merge method right after we initialize it, so all the information on how to match attribute name with respective aggregation function should be given upon initialization. Hence, __init__ should look like this:

class ResultsGroup:
    def __init__(self, lookup: dict[str, Callable]):
        self._lookup = lookup

Sidenote: `Results` are cool!

Before we get into the final look of merge, let’s remember that Results is actually a very cool class. Namely, it allows us to track which attributes we added to the object without changing the attribute access interface. In other words, we can make our results a Results object from the very beginning (namely, in _prepare method call), and then whatever got assigned, will be accessible in results.keys()! We’ll simply add line self.results = Results() in AnalysisBase._prepare.

Keep writing `merge()`

Given the coolness of the Results, now we can rely on the fact that we know for sure which attributes got assigned, and easily iterate through them. Also, since all remote_objects have the same results outline, we can pick the first one as an example, and then be sure that the rest have the same outline:

class ResultsGroup:
    def __init__(self, lookup: dict[str, Callable]):
        self._lookup = lookup
    
    def merge(self, remote_objects: list[AnalysisBase]) -> Results:
        rv = Results()

        for key in remote_objects[0].keys():
            agg_function = self._lookup.get(key, None)
            if agg_function is None:
                raise ValueError(f"No aggregation function for {key=}")
            results_of_t = [obj[key] for obj in objects]
            rv[key] = agg_function(results_of_t)
        return rv
    
    # and @staticmethod s with aggregation functions

How will `AnalysisBase.run()` look like?

Before, we had our aggregation function in _parallel_conclude method:

class AnalysisBase:
    def run(self, start, stop, step, frames, n_workers, scheduler):
        self._setup_frames(start, stop, step, frames)
        self._prepare()
        computation_groups = self._setup_computation_groups(
            start, stop, step, frames, n_workers
        )

        executor = ParallelExecutor(n_workers, scheduler)
        executor.apply(self._compute, computation_groups)

        # THIS ONE
        # --------
        self._parallel_conclude()
        # --------

        self._conclude()

now, we can do everything more explicitly:

class AnalysisBase:
    def run(self, start, stop, step, frames, n_workers, scheduler):
        self._setup_frames(start, stop, step, frames)
        self._prepare()
        computation_groups = self._setup_computation_groups(
            start, stop, step, frames, n_workers
        )

        executor = ParallelExecutor(n_workers, scheduler)
        remote_objects = executor.apply(self._compute, computation_groups)

        aggregator = ...  # will think about it later
        self.results = aggregator.merge(remote_objects)
        self._conclude()

Obviously, in order to get an appropriate ResultsGroup aggregator, we should call some method of self. Well, let’s call it exactly like this:

class AnalysisBase:
    def run(self, start, stop, step, frames, n_workers, scheduler):
        ...

        aggregator = self._get_aggregator()

        ...
    
    def _get_aggregator(self) -> ResultsGroup:
        return ResultsGroup(lookup=None)

Now, we’re running into a backwards compatibility issue – now we must have a meaningful ResultsGroup aggregator even if we’re using only backend='local', without any parallelization. In this case, however, our remote_objects is a list with 1 element, and we can simply return it in merge method:

class ResultsGroup:
    def __init__(self, lookup: dict[str, Callable]):
        self._lookup = lookup
    
    def merge(self, remote_objects: list[AnalysisBase]) -> Results:
        if len(remote_objects) == 1:
            rv = remote_objects[0]
        else:
            rv = Results()
            for key in remote_objects[0].keys():
                agg_function = self._lookup.get(key, None)
                if agg_function is None:
                    raise ValueError(f"No aggregation function for {key=}")
                results_of_t = [obj[key] for obj in objects]
                rv[key] = agg_function(results_of_t)
          return rv

note that we won’t even screw the function signature up and return the Results type still, since we’ve added self.results = Results() in _prepare().

The last touch is that not everything is stored in results – in _prepare method, attributes frames and times are also assigned, so we’ll have to merge them manually:

class AnalysisBase:
    def run(self, start, stop, step, frames, n_workers, scheduler):
        self._setup_frames(start, stop, step, frames)
        self._prepare()
        computation_groups = self._setup_computation_groups(
            start, stop, step, frames, n_workers
        )

        executor = ParallelExecutor(n_workers, scheduler)
        remote_objects = executor.apply(self._compute, computation_groups)

        # manually merge `frames` and `times`
        self.frames = np.array([obj.frames for obj in remote_objects]).sum(axis=0)
        self.times = np.array([obj.times for obj in remote_objects]).sum(axis=0)

        # apply ResultsGroup.merge()
        aggregator = self._get_aggregator()
        self.results = aggregator.merge(remote_objects)
        self._conclude()

And that’s it, a final look of AnalysisBase.run()!

How do I modify subclasses then?

For instance, let’s modify a MDAnalysis.analysis.rms.RMSD class so that it would work with our parallel backend. It has a huge _prepare method, but the only attribute that actually gets prepared is self.results.rmsd – it’s initialized with zeros of a proper shape:

self.results.rmsd = np.zeros((self.n_frames, 3 + len(self._groupselections_atoms)))

hence, in all but one remote objects value of self.results.rmsd will be zero, and we can simply add all results together to get a final result!

Let’s do that, and also add available_backends:

class RMSD(AnalysisBase):
    ...

    @classmethod
    @property
    def available_backends(cls):
        return ("local", "multiprocessing", "dask", "dask.distributed")

    def _get_aggregator(self):
        return ResultsGroup(lookup={"rmsd": ResultsGroup.ndarray_sum})

and that’s it! And the function we’ve specified here in lookup has a super simple signature – Callable[list[T], T], and all the built-in functions have literally a single line of code in them.

Conclusion

In this post, we’ve finally finished writing our AnalysisBase.run() protocol via adding a ResultsGroup class. It accepts a dict[str, Callable[list[T], T]] lookup argument, which maps attribute name to a proper aggregation function. Its merge method does exactly one job – flattens the results of all objects that got passed into it.

Altogether, our AnalysisBase.run() has changed by addition of the following methods: _compute(), _setup_computation_groups(), _get_aggregator() and available_backends(). We’ve added ParallelExecutor and ResultsGroup classes that abstract away parallel execution and results aggregation, respectively. And finally, we added multiprocessing and dask backends that are supposed to speed up the analysis!

(not anymore bi)weekly GSOC-3: writing `ParallelExecutor` class

Sun, 30 Jul 2023 00:00:00 +0000

In the previous blogpost, I briefly explained how decomposition works – we split all _single_frame() runs into independent groups that get executed in parallel. For this, we have _compute method that executes the frames group, and run method that orchestrates the _compute execution.

Here, I will explain how the actual implementation went south, and then evolved into something more complex and simple at the same time.

What was the problem?

So, the actual implementation of the run protocol turned out to be more complicated than I thought. When it was finally ready, the code looked horrible – especially the AnalysisBase.run() method, that had multiple if-else branches that would choose the exact code path depending on scheduler, and also in all this mess was the logic for the actual run(). Even though it worked (I had most of the subclasses working and passing all the tests), it would be a maintainer nightmare, and also adding new features would be close to impossible.

The mentor team noted that, and I had to come up with something that is easier to read, maintain, and add features to.

How the code looked before the change you’re talking about?

The code looked roughly like this:

def run(self, start, stop, step, frames, n_workers, scheduler):
    self._setup_frames(start, stop, step, frames)
    self._prepare()
    self._setup_bslices(start, stop, step, frames, n_workers)

    if scheduler is None:
        # we have no worker processes and only one `bslice`
        self._compute(bslice_idx=0)
    elif scheduler == "multiprocessing":
        import multiprocessing

        ...
    else:
        from dask.delayed import delayed

        computations = [
            delayed(self._compute)(bslice_idx)
            for bslice_idx in range(len(self._bslices))
        ]
        results = computations.compute()
        self._remote_results = results
        self._parallel_conclude()

    self._conclude()

And _compute method looked like this:

def _compute(self, bslice_idx):
    bslice = self._bslices[bslice_idx]
    frame_indices, frames = bslices[:, 0], bslices[:, 1]
        for idx, ts in enumerate(trajectory):
        i = frame_indices[idx]
        self._frame_index = i
        self._ts = ts
        self.frames[i] = ts.frame
        self.times[i] = ts.time
        self._single_frame()

As you can see, there is an ugly if-elif-else part of run method, and it’s not exactly clear what is happening in each of the branches (and how to test and debug it). However, all these branches are doing exactly the same thing: they apply a function to a list of computations in a parallel fashion.

What do we do?

What if we moved this functionality outside of the AnalysisBase class? Upon initialization, user would be able to configure its parallelization options, and then its apply method would execute computations with pre-configured parameters. Something like this:

class ParallelExecutor:
    def __init__(self, n_workers, scheduler):
        self.n_workers = n_workers
        self.scheduler = scheduler

    def apply(self, func, computations) -> list:
        if self.scheduler == "local":
            return [func(task) for task in computations]
        elif self.scheduler == "multiprocessing":
            import multiprocessing

            with multiprocessing.Pool(self.n_workers) as pool:
                return list(pool.map(func, computations))
        elif self.scheduler == "dask":
            import dask
            from dask.delayed import delayed

            func_d = delayed(func)
            return delayed([func_d(task) for task in computations]).compute(
                n_workers=self.n_workers
            )
        else:
            raise ValueError("wrong scheduler")

The run method would look somewhat like this:

def run(self, start, stop, step, frames, n_workers, scheduler):
    self._setup_frames(start, stop, step, frames)
    self._prepare()
    bslices = self._setup_bslices(start, stop, step, frames, n_workers)

    executor = ParallelExecutor(n_workers, scheduler)
    executor.apply(self._compute, bslices)
    
    self._parallel_conclude()
    self._conclude()

Ok, this looks much simpler! Now we can clearly see what’s going on during run, without unnecessary configuration logic.

Minor thing is that we haven’t completely got rid of the ugly if-elif-else part. Major problem is that it’s not (that) easily tested – you have to get inside a function, and then inside a particular condition. As long as we’re having conditions such as if str1 == str2, it’s not a problem, but we might have something more complicated in the future, and for this reason need to have separate functions for each computation type:

class ParallelExecutor:
    def __init__(self, n_workers, scheduler, client):
        self.n_workers = n_workers
        self.scheduler = scheduler

    def _compute_with_local(self, func, computations):
        return [func(task) for task in computations]

    def _compute_with_dask(self, func, computations):
        import dask
        from dask.delayed import delayed

        func_d = delayed(func)
        return delayed([func_d(task) for task in computations]).compute(
            n_workers=self.n_workers
        )

    def _compute_with_multiprocessing(self, func, computations):
        import multiprocessing

        with multiprocessing.Pool(self.n_workers) as pool:
            return list(pool.map(func, computations))

How do we now match it? Well, if we don’t want our code to smell but still want to map pre-configured options to the desired code path, let’s use a general mapping object – dictionary:

class ParallelExecutor:
    ...

    def apply(self, func, computations):
        computation_options = {
            self._compute_with_local: self.scheduler == "local",
            self._compute_with_dask: self.scheduler == "dask",
            self._compute_with_multiprocessing: self.scheduler == "multiprocessing",
        }
        for applicator, condition in computation_options.items():
            if condition:
                return applicator(func, computations)
        raise ValueError("wrong configuration")

Cool! No code smell, easy testing, easy usage, potential to be used in other parts of the project.

Important linguistic changes

Two paragraphs on the naming – if you’re reading this posts few months into the future, you might notice that they use different terminology from the one used in the codebase. Here’s the reason for it.

`bslices --> computation_groups`

This project was inspired by pmda and initially borrowed many things from there – namely, the bslices term, which means “balanced slices”. More specifically, it means “groups of frames that are then passed to independent worker objects”. With this explanation, it’s not exactly clear why are they slices and in which way they are balanced, so I decided to rename them into computation_groups, and respective method into AnalysisBase._setup_computation_groups().

`scheduler --> backend`

It’s nothing wrong with this term, but it turns out that dask uses the same term during standalone computations execution, or configuration of dask.distributed.Client object. In dask, scheduler can be either synchronous, processes or threads, which is vastly different from the usecases in MDAnalysis. In order to avoid confusion, scheduler is now called backend.

Conclusion

So, we figured out how to perform parallel computations in a way that a) does not disrupt the flow of the AnalysisBase.run() method; b) allows us to reuse a huge chunk of code for completely different purposes. However, we’re still not finished on how to complete our re-writing of the run method – namely, we haven’t yet figured out on how to aggregate the results from independent workers.

To be honest, this (and pytest) have been the reason of this post becoming less-than-bi-weekly, and I hope I’ll have a proper answer next time. Stay tuned!

Biweekly GSOC-2: splitting the work and writing `_setup_bslices`

Fri, 30 Jun 2023 00:00:00 +0000

In the previous blogpost, I briefly explained how I’m planning to decompose the AnalysisBase.run() method, so that its subclasses won’t notice the changes in the protocol, but at the same time will be able to perform computations in a parallel manner.

Here I’ll go through the implementation details of some methods – namely, I’ll explain how I went to a certain implementation of _setup_bslices, _compute and run.

Where are we right now?

At this moment, the AnalysisBase.run() method looks roughly like this:

def run(self, start, stop, step, frames):
    """
    Perform the calculation
    """
    self._setup_bslices(...)
    computations = []
    for bslice in self._bslices:
        start, stop, step, frames = bslice
        computations.append(delayed((self._compute)(start, stop, step, frames)))
    results = computations.compute()
    self._remote_results = results
    self._parallel_conclude()
    
    self._conclude()
    return self

Let’s think of how _setup_bslices should look like – which arguments it has and what it returns, or which attributes of self modifies?

Well, clearly, it should know about the computation limits defined earlier in the run itself. Also, it should know about the scheduling parameters, in order to be able to distribute the load more or less evenly. Also, we know that we’ll iterate over self._bslices in the run(), so we want to assign this attribute at the end. So, something like that:

def _setup_bslices(self, start, stop, step, frames, scheduler):
    n_parts = some_function_of(scheduler)
    bslices = split_evenly(start, stop, step, frames)
    self._bslices = bslices

Now, the first part with n_parts = ... is actually simple. For now, for clarity, we’ll just specify scheduler as a string that represents the scheduling backend – either additionally installed dask, built-in multiprocessing, or None representing execution in the current process, and n_parts will be n_workers – a parameter of a scheduler.

What are we actually splitting?

In order to split the work correctly, we need to know how exactly it looks like. In MDAnalysis, there are 2 mutually exclusive ways to define which frames you want to analyze:

start, stop, step – a triplet of integer numbers with semantics similar to those in range
frames – an iterable that works as a slice over trajectory. Can be of two kinds:
- frames: Iterable[int] – an iterable with numbers of frames for analysis. For example, frames = list(range(0, len(trajectory))) will be equivalent to start=0, stop=len(trajectory), step=1
- frames: Iterable[bool] – an iterable with boolean values that define which frames to take into account. For example, frames = [bool(i%2) for i in range(len(trajectory))] is equivalent to start=1, stop=len(trajectory), step=2

Now our task is to accept one of those parameters, and correctly split it into roughly equal parts, each of which will go to a separate worker. First of all, let’s simplify our task and note that accepting start, stop, step is equivalent to frames = list(range(start, stop, step)). Moreover, using frames: Iterable[bool] is actually equivalent to using frames: Iterable[int] = np.arange(len(trajectory))[frames]. Now it’s the only case we have to work with! Neat.

In order to split an iterable into n equal parts, one could start writing some iterator-based functions in python, spend some time debugging it, etc – just like I did initially. Or be slightly smarter and google an appropriate function – turns out, np.array_split does exactly what we need.

Now, the splitting function needs to first match the type of our parameters, and then apply np.array_split single split. Something like this:

def split_work_into_parts(
    n_parts: int = 1, start=None, stop=None, step=None, frames=None
):
    if any([
        opt is not None for opt in (start, stop, step)
    ]):  # using `start-stop-step` notation
        frames = np.arange(start, stop, step)
    else:  # using `frames` notation
        if is_iterable_of_booleans(frames):
            frames = np.arange(len(frames))[frames]

    frames = np.array_split(frames, n_parts)
    return frames

where is_iterable_of_booleans(frames: Iterable) looks something like:

def is_iterable_of_booleans(arr: Iterable):
    return all((isinstance(obj, bool) for obj in arr))

Unfortunately for us, we can’t simply use the result of split_work_into_parts – we not only need to split the frames evenly, but we also have to keep track of frame indices we’re using, since in _compute we’re explicitly assigning self._frame_index = i. Hence, we have to use something like enumerate before creating balanced slices, and make our function slightly more complicated:

def split_work_into_parts(
    n_parts: int = 1, start=None, stop=None, step=None, frames=None
):
    if any([
        opt is not None for opt in (start, stop, step)
    ]):  # using `start-stop-step` notation
        frames = np.arange(start, stop, step)
    else:  # using `frames` notation
        if is_iterable_of_booleans(frames):
            frames = np.arange(len(frames))[frames]

    frames = np.array(list(enumerate(frames)))
    frames = np.array_split(frames, n_parts)
    return frames

Neat! Now we can get back to writing _setup_bslices.

Writing `_setup_bslices` and using `self._bslices`

Given all functions above, the target method will now look like this:

def _setup_bslices(self, start, stop, step, frames, n_workers):
    equal_iterables = split_work_into_parts(n_parts=n_workers, start, stop, step, frames)
    self._bslices = equal_iterables

Now, we want to use these bslices in run method. We could get both indices and frames from self._bslices and then use them when creating list of computations, but since self._bslices will anyway get passed to the worker objects, let’s just use a bslice index for assigning work, hence simplifying the run code. Like this:

def run(self, start, stop, step, frames, n_workers, scheduler):
    self._setup_frames(start, stop, step, frames)
    self._prepare()
    self._setup_bslices(start, stop, step, frames, n_workers)

    if scheduler is None:
        # we have no worker processes and only one `bslice`
        self._compute(bslice_idx=0)
    else:
        from dask.delayed import delayed

        computations = [
            delayed(self._compute)(bslice_idx)
            for bslice_idx in range(len(self._bslices))
        ]
        results = computations.compute()
        self._remote_results = results
        self._parallel_conclude()

    self._conclude()

And usage of _compute will also slightly change, since we’re using bslice_idx instead of original run-inherited syntax:

def _compute(self, bslice_idx):
    bslice = self._bslices[bslice_idx]
    frame_indices, frames = bslices[:, 0], bslices[:, 1]
        for idx, ts in enumerate(trajectory):
        i = frame_indices[idx]
        self._frame_index = i
        self._ts = ts
        self.frames[i] = ts.frame
        self.times[i] = ts.time
        self._single_frame()

And that’s it! This is really the core protocol of the AnalysisBase.run() method.

Conclusion

We now know crucial part of the parallelization works – splitting work into equal parts and communicating between those parts. We haven’t yet touched on aggregation of these results – we only know that all workers, together with their results, will be stored in self._remote_results: list[AnalysisBase], but we are yet to find out how to retrieve the main AnalysisBase results from these objects.

Biweekly GSOC-1: decomposing `AnalysisBase.run()`

Fri, 16 Jun 2023 00:00:00 +0000

In the first GSOC-devoted blogpost I explained the idea behind my proposal and the approximate timeline for its execution.

Here I’ll go through the initial implementation stages – how I decided to decompose the run() method for the AnalysisBase so that I could parallelize the execution, which changes it’ll require, and how not to break all the existing code in the process. Let’s go!

`AnalysisBase` protocol

AnalysisBase is a superclass for most of the MDAnalysis objects that perform the actual analysis of trajectories. If we omit some house-keeping arguments, Its current implementation works roughly like this:

class AnalysisBase(object):
    def __init__(self, trajectory):
        """
        Initialize the run object
        """
        self._trajectory = trajectory
        self.results = ...
    
    def run(self, start, stop, step, frames):
        """
        Perform the calculation
        """
        self._setup_frames(self._trajectory, start=start, stop=stop,
                           step=step, frames=frames)
        self._prepare()

        for i, ts in enumerate(self._sliced_trajectory):
            self._frame_index = i
            self._ts = ts
            self.frames[i] = ts.frame
            self.times[i] = ts.time
            self._single_frame()

        self._conclude()
        return self
    
    def _setup_frames(self, trajectory, start, stop, step, frames):
        """
        Prepare frames that will be used for the analysis
        """
        self._sliced_trajectory = ...
        self.start = start
        self.stop = stop
        self.step = step
        self.n_frames = ...
        self.frames = ...
        self.times = ...
    
    def _single_frame(self, ...): # implemented in subclasses
        """
        Perform calculations on a single frame
        """
        do_some_computations()
    
    def _prepare(self, ...): # implemented in subclasses
        """
        Prepare the storage attributes for intermediate results
        """
        self._intermediate_data = ...
    
    def _conclude(self, ...): # implemented in subclasses
        """
        Use the intermediate results to create the final ones
        """
        self.results = some_function_of(self._intermediate_data)

Most of the computations happen in the run method – namely, here:

for i, ts in enumerate(self._sliced_trajectory):
    self._frame_index = i
    self._ts = ts
    self.frames[i] = ts.frame
    self.times[i] = ts.time
    self._single_frame()

So we must somehow parallelize the _single_frame() method, and make it run in separate processes simultaneously.

Where to apply `dask` parallelization?

The parallelization in dask, which was our framework of choice (and also was used earlier in pmda), works roughly like this:

from dask.delayed import delayed


@delayed
def simple_computation(x, y):
    do_something()


parameter_space = [(x, y) for x in range(10**6) for y in range(10**6)]
computations = delayed([simple_computation(x, y) for x, y in parameter_space])
results = computations.compute()

At the last line, dask spawns the computations among all workers – it serializes the simple_computation function and its arguments, sends them to the workers (which are either independent processes or even independent machines), does the computation and returns the results again via serialization.

Unfortunately, our case is a bit more complex than that: _single_frame uses many attributes of the AnalysisBase class itself – namely, those that are being set up before the _single_frame in the loop, and also those that each subclass might want to set up in their _prepare block. Hence, we should serialize the whole class instance, together with all its attributes. But then we can’t do _single_frame a delayed function, because then each delayed function will carry the whole serialized class instance with it, and just blow up the memory of any local or remote computer. Also, _single_frame() does some things internally, modifying the class instance, and returns None after that. So if we were to make each _single_frame() a delayed function, we’ll gather the None values in our results after that, which is kind of silly.

So, we should split the computations in parts, and submit each part to its separate worker. Within each part, we’ll have our own set of frames for computation, and then will collect all the results from these parts somewhere inside the “main” instance’s run. For the purpose of continuity with pmda, we’ll call these parts balanced slices, or bslices.

Setting up the frames for the computation in each worker would be a tedious task. Luckily, we have the function exactly for these purposes: AnalysisBase._setup_frames() does exactly this! It sets up all the necessary attributes in the class instance in a way that if we run self.run() after that, we’ll be able to successfully iterate only through frames that were prepared.

Introducing `_compute()` method

For now, each serialized instance is getting its own set of frames, bslice, and works with that. Let’s separate all this work into a separate _compute() method. How would it look like?

First, it should explicitly know the frames it was configured to work with – we’ll pass them as arguments. Second, it should configure the class instance for computations – by running _setup_frames. Third, it should run the computation loop – the one that used to be in the run method, and in order for it to run properly, we should run _prepare() first

def _compute(self, start, stop, step, frames):
    self._setup_frames(..., start, stop, step, frames)
    self._prepare()
    for i, ts in enumerate(self._sliced_trajectory):
        self._frame_index = i
        self._ts = ts
        self.frames[i] = ts.frame
        self.times[i] = ts.time
        self._single_frame()

    return self

Now, since we’re explicitly returning self here, we will have our self.results attribute, and won’t loose it while sending it to the other workers.

Decomposition or `run()`

So, run is becoming more complex – it lacks the computation loop now, but has many other methods to control the computational flow itself. The only thing left to do now is to collect all the results together from all the workers. In a way, it’s a _conclude method for the parallel part, so we’ll call it _parallel_conclude(), with the whole method now looking like this:

def run(self, start, stop, step, frames):
    """
    Perform the calculation
    """
    self._setup_bslices(...)
    computations = []
    for bslice in self._bslices:
        start, stop, step, frames = bslice
        computations.append(delayed((self._compute)(start, stop, step, frames)))
    results = computations.compute()
    self._remote_results = results
    self._parallel_conclude()
    
    self._conclude()
    return self

How should the _parallel_conclude() look like? Well, we have a list of instances of the executed AnalysisBase subclass in our _remote_results attribute. If each of them has their own results, it would look somewhat like this:

def _parallel_conclude(self):
    self.results = some_aggregation_function(self._remote_results)

Unfortunately, here we can’t avoid subclass-specific implementation – every computation is different by how it collects its intermediate results (in fact, AnalysisBase.results actually has its own Results type, which is essentially a dictionary, and can hold arbitrary data). But it’s ok, at least the run process itself is now parallelized.

Conclusion

We now have an outline for the parallel execution of the AnalysisBase.run() method, which contains proper splitting of the frames into balanced slices, and running the computation loop on each of the balanced slices. We didn’t pay much attention to the scheduling part of it, and how it’ll affect the control flow, but we’ll come back to it later in the future posts of these series!

GSOC 2023: what will I do?

Mon, 29 May 2023 00:00:00 +0000

This year I applied for the Google Summer of Code program, and luckily was selected as a contributor for MDAnalysis with my proposal for a parallel analysis project, extending the idea suggested by the developers.

In this post, I’ll briefly describe the proposal itself – motivation, technical outline, and how I expect things will change for the MDAnalysis users after the successful execution of the proposal idea.

Overview and motivation

As you might already know, MDAnalysis is a python library for analysis molecular dynamics (MD) trajectories, which is kinda clear from the start if you look at the library’s name. In turn, molecular dynamics is an ~~art~~ computational technique for simulation of molecular ensembles at various levels of detail. In structural biology, these ensembles can vary from low-level things like simulation including quantum mechanics approximations (QM/MM) to high level simulations of large systems such as a whole cell organelle or a cell itself (coarse-grained simulations).

Technically, a result of a MD simulation is a trajectory – a file containing coordinates of all simulated atoms ($10^4-10^8$ atoms) over the course of a whole simulation ($10^4-10^7$ time-steps). Producing such a file is hard and computationally intensive, and requires lots of GPU resources. Often, it’s even harder to get insights from these simulations, and requires lots of hypothesis formulation and evaluation, coupled with a data-intensive analysis of these trajectories, where libraries such as MDAnalysis come in handy.

However, currently MDAnalysis uses a single process to run an analysis of a trajectory. Despite all the efforts for the code optimization, it is hard to keep up with the size of the MD trajectories. Hence, utilizing a parallel approach for the analysis would be the next thing to do.

My proposal focuses exactly on this: I am planning to implement a parallel backend for the MDAnalysis library. The backend is planned to be implemented using dask library, allowing users to seamlessly run their analysis either on powerful local machines or various clusters, such as SLURM. There was a proof-of-concept fork of the MDAnalysis library, pmda, which implemented this idea for a subset of analysis methods implemented in the MDAnalysis. Basically, proposal idea is to refactor pmda methods into the MDAnalysis.

Technical details of the project

A key component of the MDAnalysis library is the AnalysisBase class, from which all objects that allow user to run an analysis of a trajectory are inherited. Namely, it implements a run method, that looks somewhat like that:

def run(self, start=None, stop=None, step=None, frames=None, ...):    
    self._setup_frames(self._trajectory, start=start, stop=stop, step=step, frames=frames)
    
    self._prepare()
    for i, ts in enumerate(self._sliced_trajectory, ...):
        ...
        self._single_frame()
    self._conclude()

and consists of three steps:

setting up frames for reading – may include options to analyze only a part of the trajectory in time coordinate (change time-step, start/stop, etc)
preparing for the analysis: may include preparation of some arrays storing intermediate data, etc
running analysis of a single frame
and concluding the results – e.g. average some quantity calculated for each frame separately.

For a setup with multiple worker processes, this protocol will require an additional step of first separating a trajectory into blocks. Each block will be processed with a single separate process, and also results from different blocks will potentially be concluded separately:

def run(
    self, start=None, stop=None, step=None, frames=None, scheduler: Optional[str] = None
):
    if scheduler is None:
        # fallback to the old behavior
        self._setup_frames(
            self._trajectory, start=start, stop=stop, step=step, frames=frames
        )

        self._prepare()
        for i, ts in enumerate(self._sliced_trajectory, ...):
            ...
            self._single_frame()
        self._conclude()

    else:
        self._configure_scheduler(scheduler=scheduler)
        self._setup_blocks(start=start, stop=stop, step=step, frames=frames)
        # split trajectory into blocks according to scheduler settings

        tasks = []
        for block in self._blocks:
            # create separate tasks
            # that would fall back to the old behavior
            # and schedule them as dask tasks
            subrun = self.__class__(
                start=block.start,
                stop=block.stop,
                step=block.step,
                frames=block.frames,
                scheduler=None,
            )
            dask_task = dask.delayed(subrun.run)
            tasks.append(dask_task)

        # perform dask computation
        tasks = dask.delayed(tasks)
        res = tasks.compute(**self._scheduler_params)
        self._parallel_conclude()

Which requires introducing following methods for the AnalysisBase class:

class AnalysisBase(object):
    def _configure_scheduler(self, scheduler=scheduler, **params):
        ...

    @property
    def _blocks(self):
        ...
    
    def _setup_blocks(start=start, stop=stop, step=step, frames=frames):
        # will also update `self._blocks` accordingly
        ...

    def _parallel_conclude(...):
        ...

Which is similar to the protocol implemented in pmda. Such a modular design has following advantages:

it reuses the previously existing code for prepare and conclude methods for the subclasses, and hence will successfully run on sub-trajectories
it requires developers to introduce only a proper non-default parallel_conclude method for sophisticated results combination
it allows to raise an exception for subclasses that for some reason don’t allow such parallelization (e.g. rely on results from previous frames when computing current one), via re-implementing _configure_scheduler and raising an exception in not-None cases.

Project steps

During the proposal preparation, most of the effort was put into the project planning to ensure some profit for the library even of the overall goal is unreachable within the given time-frame.

0. Proof-of-concept for the scheduler-based protocol

This part includes re-writing protocol in AnalysisBase.run() as introduced above, but using simplest backend possible – dask running on a local machine with a single job. This way, I can ensure that there are no node communication issues etc, and focus on proper storing of the values in private fields and re-factoring the protocol.

introduce _configure_backend() and _setup_blocks() in a simplest way possible – single block and one task
add simple _parallel_conclude() method
make sure that results are the same as were before on a develop branch for few test files
add tests for n_jobs=1 and test files tested manually

Here I mostly expect problems with the _setup_blocks() part, that might introduce nasty +- 1 errors in the tests. Or might not.

1. Introducing multiple jobs support

Next step would be to add complexity and expand _configure_backend() to multiple jobs, and ensure it’s still working. At this step, I’d expect serialization problems, as well as problems with reading only the necessary part of the trajectory from disk. And, similarly to the previous case, I expect some problems with the _setup_blocks().

add scheduler configuration for multiple cores
re-write tests for multiple cores
ensure the tests run with multiple cores

Also, this step would be a perfect mid-project checkpoint for several reasons:

it actually adds value to the library – if the second part of the project fails, at least the single-node parallelism would be available for users
it is complex enough to test the protocol’s viability
it allows to see the real progress even on moderate workstations

2. Introducing cluster support

This step would include adding _configure_backend() for an actual cluster including multiple nodes. Likely, this will also include introducing backend: Optional[obj] = None instead of backend: Optional[str] = None, where obj is a dask instance that can be configured prior to the actual run.

During proposal formulation, I expected most challenged here with writing tests, but thanks to the contributors team, I realized it is actually simpler than I thought, and even made up a simple toy repository dask-gh-action to write dask-based multi-processing tests.

add object-based backend=... keyword and respective configuration
add tests for the different backend options and make sure they run locally

3. Optional: adding tests for all `AnalysisBase` subclasses

It is likely that during step 2 proper pytest fixtures for dask-based execution will be written. Hence, it’s worth trying to add these fixtures to existing tests for subclasses of AnalysisBase and see if they run, and mark NotImplemented those that have valid reasons not to run in parallel.

add tests for all AnalysisBase subclasses executed with dask
mark parallel execution for those that can’t be run trivially as NotImplemented during _configure_scheduler

4. Optional: testing MDAnalysis parallel execution on a super-large MD trajectory

It would be cool to see the results on a large trajectory such as those from ANTON supercomputer listed here, for example DESRES-ANTON-11441075 (235 Gb compressed tar). This way we’ll actually show that it’s possible to analyze things that weren’t accessible before with MDAnalysis.

Fast hex to np.ndarray

Tue, 25 Apr 2023 00:00:00 +0000

Hexadecimal string lists to np.ndarray, but fast

Say you have a list of strings, representing long binary arrays (say, 2048 bits), but represented in hexadecimal format:

# compressed --> decompressed
01ab         --> [0,0,0,0, 0,0,0,1, 0,1,0,1, 0,1,1,0]
0001         --> [0,0,0,0, 0,0,0,0, 0,0,0,0, 0,0,0,0]
...          
00a0         --> [0,0,0,0, 0,0,0,0, 0,1,0,1, 0,0,0,0]

To say, a poors man’s binary file. And you want to work with these arrays using numpy. To say, you want a function:

def convert(s: str) -> np.ndarray:
    ... # do some magic

>>> s = '01ab'
>>> answ = np.array([0,0,0,0, 0,0,0,1, 0,1,0,1, 0,1,1,0])
>>> assert np.allclose(convert(s), answ)

If you have little experience and lots of patience, you’d write something like this:

def convert_1_patience(s: str) -> np.ndarray:
    rv = ""
    for letter in s:
        if letter == "0":
            rv += "0000"
        elif letter == "1":
            rv += "0001"
        elif letter == "2":
            rv += "0010"
        elif letter == "3":
            rv += "0011"
        elif letter == "4":
            rv += "0100"
        elif letter == "5":
            rv += "0101"
        elif letter == "6":
            rv += "0110"
        elif letter == "7":
            rv += "0111"
        elif letter == "8":
            rv += "1000"
        elif letter == "9":
            rv += "1001"
        elif letter == "a":
            rv += "1010"
        elif letter == "b":
            rv += "1011"
        elif letter == "c":
            rv += "1100"
        elif letter == "d":
            rv += "1101"
        elif letter == "e":
            rv += "1110"
        elif letter == "f":
            rv += "1111"
    rv_int = []
    for letter in rv_str:
        rv_int.append(int(letter))
    return np.array(rv)

Which, in more idiomatic way and usage of builtins int and bin could be re-written as:

def convert_2_builtins(s: str) -> np.ndarray:
    binarized = bin(int(s, 16))[2:].zfill(len(s) * 4)
    return np.array([int(i) for i in binarized])

So far so good! Let’s look how how it performs. It should be noted that the strings aren’t quite random – they have some domain specifics. Namely, they’re chemical fingerprints of molecules. Roughly, each bit represents whether a certain substructure is present in the molecule or not. So, I won’t be testing the functions on purely random strings, but will take some real fingerprints – 10k fingerprints, computed by wonderful chemfp package. Look like this:

$ head -n 10 my.fps
#FPS1
#num_bits=2048
#type=RDKit-Morgan/1 radius=2 fpSize=2048 useFeatures=0 useChirality=0 useBondTypes=1
#software=RDKit/2018.09.1 chemfp/1.6.1
#source=./split_0000001.txt
#date=2021-07-19T20:44:39
00420000000000000880014002000000000000000000001000000000000002000000000000000000000000000000008020000000080000100000000000000000000000000040020000000001000000000004000008108000000000000000000000000000800000000000000000000000000008000000000000000000000008000000000000000000010000000000000000000000080000000000000000000000000000000000000000000000100000000000000001100000000000000000000000000080000000000000000000000000000008000000000000024000000000000000000000000000000002000000002000000000000004008000108000000000    ZINC000000002166
00000000000000000000010000008000000000080000000000000000000000000000000020000000000000000000000020000000000000000000000000000000000000000000000010000000000000000010200000002000000000000100000000000000000000000000000000000200000000000000000000000000000000000000000002000000010000040000000000090000008000000000000800000000000008000000000000200000100000000000000000000000000000000000000000000000200000000000000000000000000800000000000000004200000080000000000000000000000002000000020000000080000100000000000000000400    ZINC000000007578
0040000000000000008000400000000000100000000000810000000000000202000000000000000000000000000000800000000008000000000000004000000000000000000000000000000000000000000400000800800000000400000040000000000080000000000a000000080000000008000000000000000000000000000000000002000000010000000000000001000000000000000000002000000000000000000000000000000000100000000000000000000000000000000000000001000080000000000000000000000000000020000000000000004000000000000000000100000000000002000000002000000000000004040000000000000000    ZINC000000031355
00000000000000000000000000000000000000000000008002000004000004000000000000000000000000000400000020000000000000400000000000000000000008000000000000000000000000000004000000000080000000000000000000000000800000000000000000000000000000c08000000000000000000000081000000000002000010000000000000001000000008000000000000000000000000000000020000000000000100080000000800000000000000100000000000000000000008000000001000000000000000000000000000020004000000000000000000100000000000002000000002000000000000000000000000000000000    ZINC000000032133

The file contains 10k fingerprints altogether – it’s a random sample of larger ZINC20 database, containing roughly 1 billion of molecules. So the reading speed of these fingerprints matters.

Let’s run the benchmark though. Load the fingerprints:

>>> with open('my.fps') as fin:
>>>     fps = [line.split()[0] for line in fin if not line.startswith('#')]

And run the benchmark. To be honest, the run on whole 10k fingerprints is terribly slow, so I cut it down to 10:

>>>  %timeit -n 10 [convert_1_patience(s) for s in fps] 
3.75 ms ± 254 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)  
>>> %timeit -n 10 [convert_2_builtins(s) for s in fps]
3.5 ms ± 102 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

Not much, isn’t it? Extrapolating to 10k fingerprints, we get horrible 3 seconds for a single file. We can do much than that! In python, when you want to speed something up for multiple similar objects, your first good take is using numpy – under the hood it uses super efficient Fortran, so the processor in turn will do all the fancy vectorization tricks with your objects.

Let’s think how our function convert would look like. There are two complicated things in the function: converting string to integer, converting integer to several 1 and 0’s, and iterating over the array. The last two steps combined could probably be done within numpy, since it’s very good with numbers. And indeed, there is a function np.unpackbits, that does exactly what’s needed – uncompresses the representation of array of numbers into single bits:

>>> a = [1,2,3]
>>> np.unpackbits(a)
TypeError: Expected an input array of unsigned byte data type

Well, that’s reasonable – for now, a contains integers of unknown size, thus numpy doesn’t know how many bits per element the resulting array should have. Let’s fix this using the smallest numpy integer type available:

>>> a = np.array([1,2,3]).astype(np.uint8)
>>> np.unpackbits(a)
array([0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0,
       1, 1], dtype=uint8)

The problem is, we now have twice as large array – we wanted 4 elements per 1 symbol, and got 8. Let’s fix this with reshaping in the final function:

def convert_3_vectorized(s: str) -> np.ndarray:
    darr = [int(l, 16) for l in s]
    darr = np.array(darr).astype(np.uint8)
    darr = np.unpackbits(darr).reshape(-1, 8)[:, 4:].flatten()
    return darr

In the last line, we simply take last 4 bits of each number, since we’re pretty sure that the first 4 bytest will be zeroes. Benchmarks:

>>> %timeit -n 10 [convert_3_vectorized(s) for s in fps]  
1.17 ms ± 108 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

Wow, that’s 3 times improvement, not bad. Let’s try a bit faster using lazy np.fromiter function instead of constructing a whole array:

def convert_4_vectorized_iter(s: str) -> np.ndarray:
    darr = (int(l, 16) for l in s)
    darr = (
        np.unpackbits(np.fromiter(darr, dtype=np.int8)).reshape(-1, 8)[:, 4:].flatten()
    )
    return darr

The timing is basically similar:

>>> %timeit -n 10 [convert_4_vectorized_iter(s) for s in fps]
1.06 ms ± 84 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

Probably, at this moment the main bottleneck lays in converting strings into integers – the builtin int function is good for arbitrary numbers, but we have a very specific domain, hexadecimal letters. In fact, this function has to work on 16 symbols only. In fact, we already wrote such a miracle in the very first function. Let’s use it:

def shortint(letter: str) -> int:
    if letter == "0":
        return 0
    elif letter == "1":
        return 1
    elif letter == "2":
        return 2
    elif letter == "3":
        return 3
    elif letter == "4":
        return 4
    elif letter == "5":
        return 5
    elif letter == "6":
        return 6
    elif letter == "7":
        return 7
    elif letter == "8":
        return 8
    elif letter == "9":
        return 9
    elif letter == "a":
        return 10
    elif letter == "b":
        return 11
    elif letter == "c":
        return 12
    elif letter == "d":
        return 13
    elif letter == "e":
        return 14
    elif letter == "f":
        return 15


def convert_5_shortint(s: str) -> np.ndarray:
    darr = (shortint(l) for l in s)
    darr = (
        np.unpackbits(np.fromiter(darr, dtype=np.uint8)).reshape(-1, 8)[:, 4:].flatten()
    )
    return darr

And the benchmark:

>>> %timeit -n 10 [convert_5_shortint(s) for s in fps]
869 µs ± 161 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

Ok, we’re surprisingly getting somewhere! Although the function is super dumb and not idiomatic, it indeed gives some boost.

Let’s work further in this direction. Let’s dream: it would be cool if we could simply use indexing instead of chain of if ... elif statements, and index some structure like return struct[letter]. Oh wait…

def shortint_dict(l: str) -> int:
    struct = "01234567890abcdef"
    struct = {
        "0": 0,
        "1": 1,
        "2": 2,
        "3": 3,
        "4": 4,
        "5": 5,
        "6": 6,
        "7": 7,
        "8": 8,
        "9": 9,
        "a": 10,
        "b": 11,
        "c": 12,
        "d": 13,
        "e": 14,
        "f": 15,
    }
    return struct[letter]


def convert_6_shortint_dict(s: str) -> np.ndarray:
    darr = (shortint_dict(l) for l in s)
    darr = (
        np.unpackbits(np.fromiter(darr, dtype=np.uint8)).reshape(-1, 8)[:, 4:].flatten()
    )
    return darr

But the benchmarks aren’t so happy about it:

>>> %timeit -n 10 [convert_6_shortint_dict(s) for s in fps]
2.48 ms ± 386 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

But wait a second: since python is dynamic, the function isn’t precompiled in any way, but run every time dynamically from top to bottom, including the dict initialization. The solution to that would be to put the dictionary somewhere where it’s computed only once. Although an experienced programmer would simply put it out of the function scope, I’d use a dirty hack here and put it as a function default parameter – it’s well known that default parameters are evaluated only once upon compilation to bytecode. So:

def shortint_dict_default_cached(
    l: str,
    struct={
        "0": 0,
        "1": 1,
        "2": 2,
        "3": 3,
        "4": 4,
        "5": 5,
        "6": 6,
        "7": 7,
        "8": 8,
        "9": 9,
        "a": 10,
        "b": 11,
        "c": 12,
        "d": 13,
        "e": 14,
        "f": 15,
    },
) -> int:
    return struct[l]


def convert_7_shortint_dict_cached(s: str) -> np.ndarray:
    darr = (shortint_dict_default_cached(l) for l in s)
    darr = (
        np.unpackbits(np.fromiter(darr, dtype=np.uint8)).reshape(-1, 8)[:, 4:].flatten()
    )
    return darr

And the benchmarks are indeed happy about it:

>>> %timeit -n 10 [convert_7_shortint_dict_cached(s) for s in fps]
936 µs ± 338 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

Ok, not bad. But we still are kind of stuck around 1 ms per 10 fingerprints, which is still too much.

Let’s agree that we’ve exsausted all the opportunities for internal letter-to-integer conversion in default python, and think about ways to do it in numpy. To do that, we need to understand what is a letter, internally. Turns out it’s quite simple: int utf-8 encoding (which we work in), all ASCII letters (meaning, not weird emojis or stuff) are represented by a single byte. In order to access these bytes, we can use a builtin function bytes or bytearray. In our case, it’s convenient to use a bytearray.fromhex constructor:

>>> s = '0123456789abcdef'
>>> bytearray.fromhex(s)
bytearray(b'\x01#Eg\x89\xab\xcd\xef')
>>> len(_), len(s)
(8, 16)

Note that the length has changed – it’s now two times shorter. That’s because the fromhex constructor wants to create a byte from every two hex letters – each of them representing 4 bits, 2*4=8 bits = 1 byte in total.

Let’s now feed these bytes into numpy. Turns out there is a convenient np.frombuffer function that accepts a buffer-like object as a first argument (for example, the result of bytearray.fromhex), and a “encoding” parameter – a numpy data type:

>>> np.frombuffer(bytearray.fromhex(s), dtype=np.uint8)
array([  1,  35,  69, 103, 137, 171, 205, 239], dtype=uint8)

Which, after bit unpacking, turns into:

>>> np.unpackbits(np.frombuffer(bytearray.fromhex(s), dtype=np.uint8)).reshape(-1,4) 
array([[0, 0, 0, 0],
       [0, 0, 0, 1],
       [0, 0, 1, 0],
       [0, 0, 1, 1],
       [0, 1, 0, 0],
       [0, 1, 0, 1],
       [0, 1, 1, 0],
       [0, 1, 1, 1],
       [1, 0, 0, 0],
       [1, 0, 0, 1],
       [1, 0, 1, 0],
       [1, 0, 1, 1],
       [1, 1, 0, 0],
       [1, 1, 0, 1],
       [1, 1, 1, 0],
       [1, 1, 1, 1]],
 dtype=uint8)

Wow, that’s the result we actually wanted! Turns out, although the resulting bytearray has halved length, a bit unpacking together with the right dtype gets the job done.

The resulting function turns out to be extremely simple then:

def convert_8_frombuffer(s: str) -> np.ndarray:
    return np.unpackbits(np.frombuffer(bytearray.fromhex(s), dtype=np.uint8))

And the benchmarks are as fast as you’d expect:

>>> %timeit -n 10 [convert_8_frombuffer(s) for s in fps]
49.2 µs ± 1.58 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

Wow, how cool is that! 20 times speed improvement by talking to the bits directly, without any conversion.

Finally, let’s remember that we always want to convert the whole array of strings, not a single string. Namely, we want to have a function that accepts a list of strings, and returns a 2D np.ndarray. To date, we used a simple list comprehension (in benchmarks) to achieve that:

def convert_9_List(fps: List[str], converter=convert_8_frombuffer) -> np.ndarray:
    return np.array([converter(fp) for fp in fps])

The benchmarking code would change a bit as well:

>>> %timeit -n 10 convert_9_List(fps)
59.5 µs ± 22.2 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

Basically, the same time.

But… isn’t the naive list comprehension something that we so successfully got rid of using numpy? It is indeed! Let’s get rid of that – concatenate all strings in the input list, and then run our fastest function on them:

def convert_9_List_join(fps: List[str], converter=convert_8_frombuffer) -> np.ndarray:
    return converter("".join(fps)).reshape(len(fps), -1)

Timing is slightly better:

>>> %timeit -n 10 convert_10_List_join(fps) 
27.6 µs ± 8.19 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

But remember, we had a very large file. Let’s try on an actual file indeed:

>>> %timeit -n 10 convert_10_List_join(all_fps)
7.06 ms ± 634 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
>>> %timeit -n 10 convert_9_List(all_fps)
31.2 ms ± 909 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

Turns out our trick got us roughly 4 times more speed. Full benchmark results:

Running convert_1_patience_L on arr length=10 in N=10
    mean time: 4478.0 us
     std time: 108.3 us
     min time: 4377.0 us
     max time: 4699.3 us
----------------------------------------
Running convert_2_builtins_L on arr length=10 in N=10
    mean time: 3396.3 us
     std time: 63.7 us
     min time: 3298.2 us
     max time: 3515.4 us
----------------------------------------
Running convert_3_vectorized_L on arr length=10 in N=10
    mean time: 1006.1 us
     std time: 11.8 us
     min time: 985.6 us
     max time: 1026.6 us
----------------------------------------
Running convert_4_vectorized_iter_L on arr length=10 in N=10
    mean time: 922.1 us
     std time: 14.7 us
     min time: 909.2 us
     max time: 959.3 us
----------------------------------------
Running convert_5_shortint_L on arr length=10 in N=10
    mean time: 609.1 us
     std time: 16.6 us
     min time: 586.8 us
     max time: 636.3 us
----------------------------------------
Running convert_6_shortint_dict_L on arr length=10 in N=10
    mean time: 2122.5 us
     std time: 53.3 us
     min time: 2025.8 us
     max time: 2242.1 us
----------------------------------------
Running convert_7_shortint_dict_cached_L on arr length=10 in N=10
    mean time: 634.0 us
     std time: 17.4 us
     min time: 608.5 us
     max time: 663.9 us
----------------------------------------
Running convert_8_frombuffer_L on arr length=10 in N=10
    mean time: 32.6 us
     std time: 3.4 us
     min time: 31.0 us
     max time: 42.6 us
----------------------------------------
Running convert_9_List on arr length=10 in N=10
    mean time: 24.2 us
     std time: 1.9 us
     min time: 23.1 us
     max time: 29.7 us
----------------------------------------
Running convert_10_List_join on arr length=10 in N=10
    mean time: 8.3 us
     std time: 1.1 us
     min time: 7.7 us
     max time: 11.5 us
----------------------------------------
Running convert_1_patience_L on arr length=10000 in N=10
    mean time: 4680224.7 us
     std time: 43320.8 us
     min time: 4601090.7 us
     max time: 4745800.8 us
----------------------------------------
Running convert_2_builtins_L on arr length=10000 in N=10
    mean time: 3977978.9 us
     std time: 608747.6 us
     min time: 3601444.1 us
     max time: 5657307.7 us
----------------------------------------
Running convert_3_vectorized_L on arr length=10000 in N=10
    mean time: 1132044.7 us
     std time: 80890.9 us
     min time: 1053527.9 us
     max time: 1334513.3 us
----------------------------------------
Running convert_4_vectorized_iter_L on arr length=10000 in N=10
    mean time: 977055.6 us
     std time: 35110.4 us
     min time: 925769.5 us
     max time: 1032194.9 us
----------------------------------------
Running convert_5_shortint_L on arr length=10000 in N=10
    mean time: 657414.1 us
     std time: 24327.4 us
     min time: 636385.1 us
     max time: 720548.2 us
----------------------------------------
Running convert_6_shortint_dict_L on arr length=10000 in N=10
    mean time: 2242470.6 us
     std time: 33085.9 us
     min time: 2186380.7 us
     max time: 2311272.5 us
----------------------------------------
Running convert_7_shortint_dict_cached_L on arr length=10000 in N=10
    mean time: 676429.0 us
     std time: 10639.6 us
     min time: 652150.7 us
     max time: 686811.7 us
----------------------------------------
Running convert_8_frombuffer_L on arr length=10000 in N=10
    mean time: 39407.0 us
     std time: 5706.7 us
     min time: 34884.7 us
     max time: 53128.1 us
----------------------------------------
Running convert_9_List on arr length=10000 in N=10
    mean time: 31637.2 us
     std time: 1273.1 us
     min time: 29342.5 us
     max time: 34028.8 us
----------------------------------------
Running convert_10_List_join on arr length=10000 in N=10
    mean time: 6841.7 us
     std time: 205.6 us
     min time: 6616.5 us
     max time: 7346.9 us
----------------------------------------

This whole thing gave us 700x performance over a naive implementation on both large files and single fingerprint strings.

Lessons learned (in order they appeared): - numpy vectorization helps - builtin functions give much shorter code - stupid functions with domain knowledge might overpower the builtins - numpy vectorization helps - to use that, you have to know the inner structure of your data - bytearrays are surprisingly nice - vectorization over vectorization helps

Structural insights into the effects of glycerol on ligand binding to cytochrome P450

Sun, 01 Jan 2023 00:00:00 +0000

📄 Acta Crystallographica Section D Structural Biology (2023)

Neat piece of evidence showing that glycerol (a common cryoprotectant in crystallography) can significantly effect the binding pose, constrain it, and make the conformational ensemble of protein-ligand compless less rich.

Structural basis for receptor selectivity and inverse agonism in S1P5 receptors

Fri, 12 Aug 2022 00:00:00 +0000

Article link

The story behind this paper is relatively simple – there’s a receptor, there’s a ligand, let’s make a crystal structure of their complex, and get some insights on how the receptor works. The physiological relevance (tm) please find in the paper, while here I’ll go through some technical challenges that are more important to me.

Crystallography and computer vision

A technical issue with getting the structure was very high upstream scattering background on the images. Namely, while a normal dataset contains has background at around 15 detector counts, and signal (=diffraction peak) roughly above 50, this dataset had the same diffraction peak strength while background was around 2000 counts. I’ve tried many different things while trying to get rid of that, and also learned approaches like NMF for denoising. Luckily, the background was smooth enough that the sliding-window noise subtraction actually helped a lot.

Alphafold boom

Basically, Alphafold 2.0 appeared when we were writing the paper, and all structural biologists across the world were scared they’ll loose their jobs (spoiler: that kinda didn’t happen). Fun thing we did was that we benchmarked our structure in a mock ligand screening study, and showed that it’s actually much more useful than all Alphafold models (or, more precisely, the best of ensemble of 50 different models), and that this ensemble performs basically at a random baseline level. Structures of other receptors were more useful, i.e. had non-random ROC AUC and enrichment factor, while our structure was better by quite a bit.

Google patents

The last fun thing was that although we had our ligand, we didn’t know anything about similar ligands, even though we suspected it arose from a medicinal chemistry study. I found out that you can actually search Google Patents by SMILES string, and by entering ligand’s SMILES there, found that there was a benchmark of similar ligands by ONO Pharmaceuticals! We took that published benchmark as a reference, and showed that our structure is quite good at dissecting binders from non-binders, within this particular group of ligands at least.

Structural basis for receptor selectivity and inverse agonism in S1P5 receptors

Fri, 12 Aug 2022 00:00:00 +0000

📄 Nature Communications (2022)

Super cool paper, many challenging tasks – patent parsing, docking, serial crystallography, computer vision, etc. For more detail, please read the article description.

Structural aspects of photopharmacology: insight into the binding of photoswitchable and photocaged inhibitors to the glutamate transporter homologue

Fri, 15 Jan 2021 00:00:00 +0000

📄 Journal of the American Chemical Society (2021)

Photo. Switchable. Ligands. A prospect towards time-resolved studies and other things. Not by me though, I only did docking to confirm relatively unclear binding poses revealed by crystallography.

Structure-Based Virtual Screening of Ultra-Large Library Yields Potent Antagonists for a Lipid GPCR

Thu, 03 Dec 2020 00:00:00 +0000

📄 Biomolecules (2020)

Large-scale VLS that didn’t work as good as we hoped, but gave (at least me) a good perspective on such studies.

Small-wedge synchrotron and serial XFEL datasets for Cysteinyl leukotriene GPCRs

Thu, 12 Nov 2020 00:00:00 +0000

📄 Scientific Data (2020)

Huge dataset with serial crystallography data across wide range of resolutions and techniques.

Molecular mechanism of light-driven sodium pumping

Fri, 01 May 2020 00:00:00 +0000

📄 Nature Communications (2020)

Very cool study that shows a serial crystallography (with a jet!) dataset collected at ALBA synchrotron. The structure itself reveals a “trapped” state of KR2 rhodopsin.

Structural basis of ligand selectivity and disease mutations in cysteinyl leukotriene receptors

Fri, 06 Dec 2019 00:00:00 +0000

📄 Nature Communications (2019)

CysLT2 paper, with lots of docking and crystallography.

Structure-based mechanism of cysteinyl leukotriene receptor inhibition by antiasthmatic drugs

Fri, 11 Oct 2019 00:00:00 +0000

📄 Science Advances (2019)

CysLT1 paper, with lots of molecular dynamics, functional tests, and hopes.

Structural insights into ion conduction by channelrhodopsin 2

Fri, 24 Nov 2017 00:00:00 +0000

📄 Science (2017)

Very long awaited publication, which was also a very random co-authorship for me, though I’m very grateful for that.

Element-specific density profiles in interacting biomembrane models

Thu, 09 Feb 2017 00:00:00 +0000

📄 Journal of Physics D: Applied Physics (2017)

My very first publication, written during DESY Summerschool program. I validated someone else’s data processing, and got the same results, confirming that it’s a solid data processing pipeline.

marinegor.dev

marimo: Python notebooks that are just Python

Why it’s cool

A live example

Split-and-combine parallelization for MDAnalysis AnalysisBase subclasses

Secondary structure analysis with DSSP for MDAnalysis

A subgroup of light-driven sodium pumps with an additional Schiff base counterion

On the methods of studying protein-ligand interaction dynamics

Regression-based Accelerated Virtual Screening

Regression-Based Active Learning for Accessible Acceleration of Ultra-Large Library Docking

Custom Design of a Humidifier Chamber for In Meso Crystallization

MDAnalysis UGM 2023: MDAnalysis Parallelization @ GSOC

Parallelization of MDAnalysis

Motivation

Goals

Sneak peak: results

Sneak peak: results

Implementation

Implementation

Implementation

Implementation

Implementation

Implementation

Implementation

Implementation

Implementation

Implementation

Add to your subclass

How fast

How fast

When should I use it

When should I use it

Future features

Future features

Acknowledgements

Acknowledgements

Taskfile + uv for running python tests

Glossary

What is a Taskfile

What is uv

Taskfile example: MDAnalysis

GSOC-2023: summary

A short description of the goals of the project.

What you did.

The current state.

What’s left to do.

What code got merged (or not) upstream.

How others can use it

Conclusion

(again bi)weekly GSOC-4: finally about results aggretation

What are we trying to do?

First ideas

Why first ideas are bad

Good ideas: another new class

Sidenote: Results are cool!

Keep writing merge()

How will AnalysisBase.run() look like?

How do I modify subclasses then?

Conclusion

(not anymore bi)weekly GSOC-3: writing `ParallelExecutor` class

What was the problem?

How the code looked before the change you’re talking about?

What do we do?

Important linguistic changes

bslices --> computation_groups

scheduler --> backend

Conclusion

Biweekly GSOC-2: splitting the work and writing `_setup_bslices`

Where are we right now?

What are we actually splitting?

Writing _setup_bslices and using self._bslices

Conclusion

Biweekly GSOC-1: decomposing `AnalysisBase.run()`

AnalysisBase protocol

Where to apply dask parallelization?

Introducing _compute() method

Decomposition or run()

Conclusion

GSOC 2023: what will I do?

Overview and motivation

What is `uv`

Sidenote: `Results` are cool!

Keep writing `merge()`

How will `AnalysisBase.run()` look like?

`bslices --> computation_groups`

`scheduler --> backend`

Writing `_setup_bslices` and using `self._bslices`

`AnalysisBase` protocol

Where to apply `dask` parallelization?

Introducing `_compute()` method

Decomposition or `run()`

3. Optional: adding tests for all `AnalysisBase` subclasses