ZnTrack: Data as Code

Abstract.

The past decade has seen tremendous breakthroughs in computation and there is no indication that this will slow any time soon. Machine learning, large-scale computing resources, and increased industry focus have resulted in rising investments in computer-driven solutions for data management, simulations, and model generation. However, with this growth in computation has come an even larger expansion of data and with it, complexity in data storage, sharing, and tracking. In this work, we introduce ZnTrack, a Python-driven data versioning tool. ZnTrack builds upon established version control systems to provide a user-friendly and easy-to-use interface for tracking parameters in experiments, designing workflows, and storing and sharing data. From this ability to reduce large datasets to a simple Python script emerges the concept of Data as Code, a core component of the work presented here and an undoubtedly important concept as the age of computation continues to evolve. ZnTrack offers an open-source, FAIR data compatible Python package to enable users to harness these concepts of the future.

1. Introduction

Large-scale computing is steadily becoming the norm in many industrial and research settings. With the rise of large-scale compute clusters, computational sciences have risen equally as fast in the size and number of experiments being performed on the newly available resources (Shalf, 2020). This has coincided with extraordinary advances in machine learning which is now ubiquitous in science and industry applications. Nevertheless, the challenge doesn’t end with the capability to implement extensive simulations or machine learning models. With the growing number of simulations and models comes an increase in parameters and workflows that must be tracked and stored efficiently, not only for reproducibility but also for distribution. Furthermore, once a workflow is applied in production, it is often desirable to submit the results to a publication or, more often, share models with other users in one’s community. The challenges in managing this generated data remain an issue in many fields (Allison et al., 2016; Anderson et al., 2007; Peng, 2011; Stoddart, 2016; Klump et al., 2021). Sharing data in a standardized format can be challenging because there can either be many different formats to choose from or none that fulfill all the requirements of new research questions. New formats can be introduced or generalizations attempted, but naming conventions and standardized data formats have a finite degree of complexity that they can handle before becoming overcrowded and unusable (González et al., 2007; Oliveira et al., 2020). When discussing data, it is often overlooked that most of the time software was used to generate it in the first place. This is to say that a simulation script has been written, a configuration file parsed, or command line interface (CLI) arguments have been used to initiate this generation of data. In this case, it is reasonable to say that all information required to reproduce the data is contained in the code used in its generation. Therefore, it is a convenient idea to provide a simple interface for sharing this code, along with its results, as data, i.e., to construct an interface capable of sharing Data as Code (DaC).

In this work, we present ZnTrack, a Python package designed to address these requirements. The use of the Python programming language aligns with its popularity in the scientific community, especially in data science and machine learning (JetBrains, 2022). ZnTrack is built on the idea that once code has been written to generate data, no matter how complex the workflow, this is all that is required to reproduce or share the data. Central to this goal is ZnTrack’s construction on top of the Data Version Control (DVC) framework, which provides a convenient interface to treat versions of code, i.e experiments, as commits or tags alongside potentially large data files inside a single repository driven by the GIT version control system. ZnTrack allows users to store their data alongside the code used to generate it. The usage of a data remote, together with well-established version control infrastructure such as GitHub, GitLab, or Bitbucket, is automatically managed by utilizing DVC. With ZnTrack, we combine the universal applicability of GIT and DVC with a dynamic and flexible interface driven by the Python programming language.

The purpose of this paper is to introduce the ZnTrack software to the community both on a technical level as well as through practical examples. Initially, we introduce the concepts on which ZnTrack is built, these being computational workflow design on graphs, version control with GIT, and the DaC paradigm. Following this overview of the theoretical aspects of the package, the architecture of ZnTrack is presented and discussed along with special mention of certain key features. Finally, two use cases are explored in order to demonstrate the applicability and strengths of this new technology. We showcase how ZnTrack can be used for purely Python-based workflows as well as how it can be expanded to work with other software.

2. Related Work

2.1. Workflow and Data Management

The concept of workflow and data management has been of interest for years. Research groups and organizations have invested in the field on both theoretical and product-driven levels. The result of this investment has been the emergence of general frameworks such as Apache Airflow (Haines, 2022), kedro (Alam et al., 2023), snakemake (Mölder et al., 2021), luigi (Rieger et al., 2017) and others (Fitschen et al., 2019; Luo et al., 2021; Adorf et al., 2018; George and Saha, 2022; Amstutz et al., 2022; Dask Development Team, 2016) which are widely used today. Most of these approaches target large and complex workflows and are for organizations comprising many researchers. In most cases, a specialized setup process is required to interface with or construct a new form of these services. Additionally, the majority of these pre-existing solutions depend heavily on databases hosted locally or on servers. This introduces a significant degree of overhead for new users and typically requires experts in data structures and software engineering for maintenance. Such problems might help to explain why many academic research groups have not adopted these strategies thus far (da Silva et al., 2021; Alam and Roy, 2022). Several groups have realized these shortcomings and developed tools to fill the needs left behind. These include MLFlow (Chen et al., 2020) or wandb.ai (Biewald, 2020) which are now widely known in the academic community and are even integrated into well-known software packages, for example, MACE (Batatia et al., 2022) and NequIP (Batzner et al., 2022) from the quantum chemistry community. In many cases, modern research software groups go to the trouble of integrating their own data management systems into their code as in pymatgen (Ong et al., 2013), MDSuite (Tovey et al., 2023) or pyiron (Janssen et al., 2019). However, even in these cases, what is missing is a unified infrastructure for the storage and sharing of data. To this end, products such as DVC (Castro, 2023) have emerged to combine parameter and workflow tracking with established version control tools like GIT.

DVC enables the use of version control tools like GIT for managing large data files. Additionally, it provides a comprehensive workflow management system while remaining compatible with the entire GIT ecosystem. DVC primarily utilizes a CLI and YAML configuration files, making it universally applicable to all file formats and the tools used to generate them. However, this requires users to adapt their code to interface with DVC, which leads to an overhead and might limit code flexibility. To address some of these limitations, DVC offers a Python API for accessing data, although it does not currently provide a public Python API for constructing or distributing workflows.

2.2. Python Interfaces

For the DaC concept, we want to take a look at existing datasets that can be accessed through Python packages. This method of accessing datasets was already available in the early versions of TensorFlow (Abadi et al., 2015) and PyTorch (Paszke et al., 2019). An example of this is the MNIST (Lecun et al., 1998a) dataset, which is made available among other packages through torchvision and can be downloaded through a simple API call.

⬇

1from torchvision import datasets

2trainset = datasets.MNIST()

This concept is further improved by two of the most prominent tools for sharing training data or machine learning models, namely huggingface (Wolf et al., 2020) and kaggle (Kaggle, 2023), both of which provide a simple interface to download and use datasets.

⬇

1from datasets import load_dataset

2dataset = load_dataset("mnist")

They provide not only a simple-to-use Python interface but also a sophisticated website that enables users to search for datasets hosted along with detailed descriptions of the available data, as well as different versions of the datasets if they are available. Beneficial to these tools is the accessibility of metadata and the searchability provided by the website. However, their benefits are limited to a specific set of data made available through their platforms and they do not provide a general solution to all the challenges described above.

2.3. Code-Driven Developments

Datasets or models made available through code can be more easily integrated into workflows. To fully benefit from the code interface, it is necessary to transition from static configuration files to dynamic script-based interfaces. Due to its low barrier of entry and ubiquity in data sciences, the Python programming language is the ideal tool for this task. One successful example of the transition from static configuration files to dynamic script-based interfaces is Infractsture as Code (IaC) (Howard, 2022). IaC defines computing resources through code, rather than relying on specialized configuration files. This approach provides greater flexibility and allows for the utilization of programming language features, such as loops and conditionals. The scripts can be version-controlled in the same way as traditional configuration files, while being more transferable, allowing for better scaling, and typically being easier to maintain. By automating parts that would be redundant in static configuration files, IaC also reduces the risk of misconfigurations. Finally, the code can be documented in such a way as to make it easier to understand and maintain. This flexibility is already available in workflow managers such as Airflow but is lacking in the data management tools mentioned above. Bringing these concepts together can be seen as a prerequisite for DaC.

3. Theory

With this overview of related work, we want to introduce the DaC paradigm, showcase how it compares to existing solutions and highlight key differences. Thereby putting the focus on workflow and data management as well as collaboration.

3.1. Version Control with GIT

Version control systems are essential tools in software development, allowing developers to track changes to source code and collaborate with other team members effectively. GIT is a widely used distributed version control system that has gained popularity due to its efficiency, flexibility, and ease of use. GIT uses a decentralized approach, allowing each developer to maintain their local repository and then merge changes with another repository when ready. This approach provides several benefits, including faster processing times, improved collaboration, and the ability to work offline without requiring a connection to a central server. One of the key features of GIT is its support for branching and merging. Developers can create multiple branches of the same codebase, allowing for experimentation, feature development, and bug fixing without affecting the main branch. GIT is compatible with various operating systems, making it easy for developers to work on different platforms. Changes to the codebase are stored in GIT commits. Each commit is assigned a unique identifier based on the changes made to the code. This makes GIT ideal for human readable file formats. On the contrary, large files, especially if they are compressed and a small change to the content will update the entire file, are unsuitable to be tracked with GIT.

3.2. Data Version Control

DVC (Castro, 2023) can be employed to address the challenges posed by large data files in GIT repositories.

It achieves data versioning by computing a hash value for each file, and only this hash value is versioned using GIT. When requested, the hash value can be utilized to retrieve the data associated with a specific commit from the data storage. The storage options include local directories as well as remote storage such as object storage, WebDAV, and others (refer to Figure 1).

Using GIT in conjunction with DVC not only enables versioning of data but also promotes better collaboration and reproducibility. These benefits make it evident why this combination is advantageous for research.

Figure 1. Combination of GIT and DVC in a local repository connected to a GIT remote and DVC data remote.

Version controlling both data and code in a GIT repository facilitates easy experimentation. By versioning input parameters, code, and outputs, experiments can be defined and stored as commits. Figure 2 illustrates how experiments can be performed in such a way. The GIT and DVC CLIs enable comparison of experiments. The best experiment can be promoted to a commit on a new branch, while changes to model architecture, training data, and parameters are tracked within a single repository. This streamlined approach enhances the efficiency of experimentation and development processes.

Figure 2. Experiment versioning using GIT. Each Experiment represents a detached commit. The best experiment is committed and new experiments are performed based on this commit.

Experiments can be compared using the DVC CLI as well as graphical user interfaces (GUIs) such as the Visual Studio Code DVC extension. The corresponding data files can be accessed through a fsspec (python, 2023)-compliant Python API, that supports local and remote storage.

3.3. Computational Graphs

A core concept of modern computation is the representation of a workflow as a computational graph. Within a computational graph, workflows are defined through nodes $N$ that represent variables or operations involved in a computation. The flow of information and the dependencies between nodes are represented by the computational graph’s edges $E$. Therefore, each connected node acts as a function on its inputs and passes its resulting outputs to its successors, as defined by the edges $E$ (Owhadi, 2022; Tinhofer et al., 2012).

(1) $$\begin{split}\text{Let }G&=(N,E)\text{ be a graph with a node set }\\ N&=\{1,2,3,4,5\}\text{ and an edge set }\\ E&=\{(1,2),(1,4),(2,3),(4,5),(5,3)\}.\end{split}$$

Each tuple in the edge set $E$ is a directed connection from one node to another. An edge $(1,2)$ from node 1 to node 2 indicates that the output of node 1 is used as an input to node 2. In other words, node 2 depends on the output of node 1.

Figure 3. Illustration of a DAG.

In the context of data-driven workflows, computational graphs are a powerful tool for managing dependencies, parameters and results. We will focus on DAGs for the remainder of this work as illustrated in Figure 3. Connections within a DAG are directional, i.e. the information is only passed from the preceding node to the following one as indicated by the tuples in the edge set $E$. Furthermore, it is prohibited to form cycles or loops (Thulasiraman and Swamy, 1992).

Such graphs are broadly applied in machine learning and data analysis (Lecun et al., 1998b; Schulman et al., 2015). Graphs allow for the identification of regions that are independent of each other and can be embarrassingly parallelized or optimized together (Herlihy et al., 2020; Sabne, 2020). They can be used in the context of automatic differentiation to compute gradients and are the basis of many machine learning frameworks such as TensorFlow (Abadi et al., 2015), PyTorch (Paszke et al., 2019) or JAX (Bradbury et al., 2018).

Computational graphs can be elegantly described in makefile-like formats, where the inputs, outputs, and corresponding functions are explicitly defined for each node. Within the context of DVC, a comprehensive toolset is provided to generate and persist these graphs in the YAML file format. DVC leverages the concept of node-specific hash values, computed based on the inputs and outputs associated with each node. This approach enables DVC to efficiently determine when recomputation of results is necessary or if they can be loaded from a cache based on previous runs.

3.4. Data as Code

A key challenge in cooperating on data driven projects is the availability, accessibility and documentation of data sources. Given a data repository that provides easy access to the data, a collaborator is still required to understand the chosen data format. For making changes to the data, even more in depth knowledge about the construction of the data is required and essential parts of the computational workflow are often not shared alongside the data. In the ideal scenario the data itself contains all of this information. With DaC, we build on the idea that a single entity is responsible for the generation, storage and interface of data. It should be possible to version and share this entity with collaborators. Such a framework is illustrated in Figure 4.

Figure 4. Illustration of the DaC paradigm.

The idea is to provide access to data by sharing the code that was used to generate it, thus rendering the problem of attached meta-data by an arbitrary format redundant as it is an integral component of the code itself. By combining data generation and interfacing within the same entity, we ensure that the data remains readily accessible. Additionally, documentation is consolidated in a single location for both data generation and interface. Nevertheless, compatibility can be ensured by abstraction of the interface and adherence to existing ontologies and code formats that are only altered if not sufficient for the given task.

In this code-centered data paradigm, multiple single-responsibility entities can be assembled to construct a computational graph, facilitating workflow management. These entities serve as the nodes on the computational graph. To enhance usability, when a user runs the shared code, cached files are loaded and made available instead of re-executing the data generation process. This approach promotes both workflow management and reproducibility, making them core features of the DaC concept. Drawing inspiration from the transformative impact of version control on software development (Brindescu et al., 2014), DaC provides a viable pathway towards achieving findable, accessible, interoperable, and re-usable (FAIR) data standards.

3.5. FAIR Data and Code Style

For successful collaboration in data-driven research, it is important to agree on a set of principles that guide the development. Introduced in 2016, the importance of FAIR data (Wilkinson et al., 2016) is more relevant than ever. With the increasing amount of data that is generated, it is important to make data accessible and reusable. In addition to the application of the DaC principles to achieve FAIR data, good code practices are also essential. Using an object-oriented approach, the SOLID principles are a good starting point for high-quality code design. These principles recommend designing software in a way that is easy to maintain and extend (Martin, 2017). The SOLID principles refer to a set of guidelines, namely single responsibility (Martin, 2003), open-closed (Meyer, 1988), Liskov substitution (Liskov, 1987), interface segregation (Martin, 2003) and dependency inversion (Martin, 2003). In addition to these theoretical code design aspects, the use of opinionated code formatters and linters, along with detailed in-code documentation, can enhance collaboration.

4. Architecture

In this section, we will introduce the architecture of the ZnTrack package. We will highlight the concept of DaC and its implementation in ZnTrack. Furthermore, we will describe the different features of the package in more detail and provide exemplary code snippets.

4.1. ZnTrack

To realize the DaC paradigm, data must be interfaced through code. For this, an interface through a programming language is needed. DVC and GIT, while useful for data and code versioning respectively, provide mostly CLI tools to generate a DaC infrastructure but do not fulfill all requirements to be considered DaC on their own.

ZnTrack builds on the synergy between GIT and DVC to create a cohesive framework for establishing a unified repository where data and code coexist. Embracing the DaC paradigm, ZnTrack empowers developers to effortlessly integrate data interfaces and code components in Python. Through the definition of abstract dependency attributes and outputs, the code facilitates type checking and serves as a comprehensive documentation resource for both data and code.

Additional benefits include the ability to easily share code and data with collaborators through platforms such as GitHub or GitLab. These platforms provide further infrastructure for managing issues, community discussions and code review. Furthermore, whilst most current workflow management systems require a database, using GIT as the underlying framework allows for serverless and distributed workflow management. This simplifies the migration of existing projects to a DaC workflow. Users benefit doubly from this, as setup and maintenance efforts are drastically decreased.

The core components of the ZnTrack package are split according to the components of a computational graph:

Nodes:

The Node base class is the single interface to the code and the data. All data generation and data access are handled through this class. This includes not only the input data but also parameters, metrics, plots and the produced output data. Having a single interface responsible for one task strengthens the SOLID principles described in Section 3.5. Furthermore, it allows for documentation of the code and data in a single place. Data can be tested besides the code using common testing frameworks.

Edges:

The Node defines the expected dependency types but the connections between the Nodes are handled by the ZnTrack Project. The Project acts as the interface to the computational graph. It handles the GIT and DVC commands and gives access to experiment handling. Through the Project interface, Nodes can also be organized in groups for a better overview.

To set up a DaC Project with ZnTrack, only the initialization of a repository with GIT and DVC is required. This server-less setup makes it easy to create new projects with minimal overhead. Later on, the projects can be pushed to a server to share the results with collaborators.

Structuring experiments in a very specific way can come with an overhead of time and might not be easily adapted by scientists from different areas. To ensure a user-friendly introduction, ZnTrack provides all necessary tools to create a DaC infrastructure through a Python interface. Minimal knowledge is sufficient to track your experiments using GIT and DVC in this way. Additionally, ZnTrack enables the development of new tools that inherently incorporate FAIR data management, parameter tracking, and distributed computing.

4.2. Defining Nodes

As described in Section 3, a Node is a single step in a computational workflow. It is defined by a set of parameters and a function that is executed when the Node is run.

At its core, a Node can be described as a single function on the highest level control flow, i.e. what we from now on call the computational graph. In the ZnTrack package, Nodes can be defined as Python functions or classes. In this work, our primary focus will be on the class-based approach. The main advantage of using classes, as opposed to functions, is that they are stateful in this context, making it easier to access results.

By using classes to define Nodes, we can store the parameters of the Node as class attributes. This capability allows us to conveniently access these parameters later on and establish a direct connection between the parameters, the data, and the function that is executed.

The Node class features an abstract run method, which is intended to be overridden by the user. This allows the user to define the specific functionality that should be executed when the Node is run.

To facilitate ease of use, parameters, dependencies, and outputs are defined as class attributes. These attributes are leveraged to automatically generate the class constructor. This automation simplifies the process of initializing and configuring Node instances. It also allows for class inheritance and abstract classes, which can be used to define a common interface for a group of Nodes. The inputs and outputs to a Node that are available in this way are summarized in Figure 5. We differentiate between manual data serialisation (MDS) and automatic data serialisation (ADS). Attributes suffixed with _path are MDS attributes, i.e. they are file paths. The way data is stored at these locations is entirely up to the user. In contrast, attributes without the _path suffix are ADS attributes. These attributes are serialized by the ZnTrack package and provide a convenient way to store and access data.

Figure 5. The inputs and outputs of a node are split into MDS and ADS attributes. The MDS attributes describe file paths whilst the ADS attributes can contain arbitrary data and are managed by ZnTrack.

Managing data from different commits or streaming data from the data remote is built into the ADS attributes, whilst for MDS attributes the user must make use of the Node.state.fs filesystem interface. The Node can then be deserialized based on any commit in the repository. Alongside the Node all data from the specific commit will also be available. With the availability of the code, a Node can also be deserialized anywhere by supplying it with a link to the repository.

⬇

1import zntrack

2

3class MyNode(zntrack.Node):

4    data: str = zntrack.deps_path()

5    result: float = zntrack.outs()

6    shift: float = zntrack.params()

7

8    def run(self) -> None:

9        with open(self.data, "r") as f:

10            self.result = float(f.read()) + self.shift

A simple example can be illustrated with the class MyNode. The Node takes data as an input file and adds a parameter called shift to the data value read from the file. Finally, the Node serializes the output variable result of the computation to disk. It is now possible to use MyNode as a normal Python class and test its functionality.

4.3. Defining Graphs

To make use of ZnTrack and DVC features, we need to create a graph and add an instance of MyNode to it. Therefore, we instantiate a project, which manages the graph. Within the project’s context manager, connections between Nodes are defined and parameters are set. This code will only define the graph but not execute it. All this information will be serialized and stored in human-readable configuration files that can be tracked via GIT.

⬇

1import zntrack

2

3with zntrack.Project() as project:

4    node1 = MyNode(data="data.txt", shift=1.0)

5    node2 = MyNode2(value=node1.result)

6

7project.build()

4.4. Parallelization, Deployment and Interaction

To execute the graph built within zntrack.Project, multiple options are available. All of these options share the common feature that a single command is all that’s needed to assess the entire graph and execute only the necessary parts, i.e. only run the Nodes with changed dependencies that haven’t been executed yet. This design essentially turns each Node in the graph into a checkpoint, facilitating straightforward debugging and rerunning.

The DVC package only allows for sequential execution of the computational graph. As described in Section 3.3 the computational graph can be used to parallelize the execution of Nodes that are capable of running in parallel. Therefore, we developed an additional graph executor that uses the Dask (Dask Development Team, 2016) package to parallelize the execution of the DVC workflow. Our dask4dvc package enables the graph to be run in parallel and allows us to easily deploy the graph on a cluster, using the distributed package. This allows for efficient utilization of computational resources when executing the graph.

Parallelization of a single graph is often not the only requirement for a scientific workflow. In many scenarios, the same workflow must be executed with different parameters. The concept of running the graph with different parameters is called an experiment in DVC. In DVC, experiments are run using a queueing system based on Celery and Hydra.

⬇

1dvc exp run --queue -S MyNode.shift=1,2,3,4

2dvc exp run --queue -S MyNode.shift=range(5, 10)

The code above shows how to queue experiments with different parameters. If more flexibility is required, ZnTrack enables experimentation through Python, allowing for greater control and customization in the process.

⬇

1for x in range(5):

2    with project.create_experiment() as exp:

3        my_node.shift = x

The experiments can be executed from within Python or using one of the aforementioned graph executors. This approach can also be combined with parameter optimization libraries such as Optuna (Akiba et al., 2019) or Ray Tune (Liaw et al., 2018) to automate the process of finding optimal parameters whilst keeping all results and avoiding redundant computations.

Another powerful tool that can be brought to different branches of scientific research is the automatic deployment of predefined tasks such as model training or data analysis via continuous integration (CI) setups (Iterative, 2023a). Changes in a parameter file can automatically trigger e.g. training of a new machine learning model using a CI setup. In this way, complex pipelines can technically be initiated with a click of a button even from mobile devices. Another example would be the automatic analysis of experimental results through a CI pipeline. If a screening through many experiments with recurring analysis is required this can be automated and parallelized easily.

Finally, all of these steps are inherently fully reproducible, because all of the data is present in the repository or connected data storage. This makes validating experiments much easier and fulfills the FAIR principles whilst simplifying the process of storing experiment parameters and results in the same way. This can even be achieved on a multitude of computing hardware, utilizing containerization on Nodes.

4.5. Analyzing Results

A key task when running multiple experiments is the analysis and comparison of results. There are multiple ways to achieve this. The foremost is to use the DVC CLI to compare the results of different experiments. This allows us to easily compare predefined metrics and plots.

⬇

1dvc metrics show

2dvc metrics diff

A more convenient way to gather results is by loading the full Node instance into a Python kernel. This will make the parameters in combination with the results available through the Node attributes. To avoid loading not requested data into memory, all attributes are lazily evaluated. This means that the data is only loaded when the attribute is accessed.

⬇

1my_node.load()

If the instance is not present or multiple versions should be compared, one can also load the results from a specific revision and remote. Possible revisions are the commit hash or a tag, as well as DVC experiment names.

⬇

1my_node = zntrack.from_rev(

2    "MyNode",

3    rev="v2",

4    remote="https://github.com/user/repo"

5)

With this approach, arbitrary analysis can be performed on the data. Ultimately, an analysis node can be written with multiple nodes - even from different revisions or remotes - as dependencies and the analysis can be performed on the graph level.

4.6. Collaboration

Providing results to the scientific community and collaborating during the development of a project is also facilitated by the DaC paradigm. Although all of the described principles can be used offline, their full potential is only realized when used in a collaborative environment. DaC can be realized by sharing repositories on a platform like GitHub or GitLab. Furthermore, specialized platforms such as Iterative Studio (Iterative, 2023b) or DagsHub (dagshub, 2023) allow more streamlined access to data and provide graphical user interfaces to run and compare results. All of these tools are compatible with ZnTrack.

Additionally, Nodes written with ZnTrack can be shared as Python packages and can be made available, e.g. through PyPi, building a FAIR data and DaC infrastructure. This allows others to access models and data easily. As data and models are stored together, they are automatically compatible and can easily be used and adapted by others in their own projects.

This becomes even more important when considering that many new software applications will rely on data-driven methods and often require the inclusion of machine learning models. These models can be too large to be effectively managed through GIT version control but still necessitate seamless integration into the codebase. By adhering to DaC principles, one can avoid the need to version control the model and code separately, thus mitigating the risk of incompatibilities.

5. Demonstration

5.1. Atomistic Simulation

The ZnTrack package was developed with a strong focus on atomistic simulations. These often require experimentation and can be computationally expensive. Therefore, keeping track of all the experiments and their results is crucial. In this section, we will demonstrate how ZnTrack can be used to manage and deploy computational graphs for such simulations. Furthermore, we will introduce new ways of sharing simulation data and making results accessible to and reproducible by other researchers.

5.1.1. Molecular Dynamics Simulation

To illustrate the process, a molecular dynamics (MD) simulation for a system of molten sodium chloride is deployed. The example’s workflow consists of generating a 3D structure from SMILES (Weininger, 1988), generating a simulation box and running the actual simulation. We use the rdkit package (Landrum et al., 2023) together with packmol (Martínez et al., 2009) to generate the structure and Lammps (Thompson et al., 2022) to run the simulation.

We run multiple experiments for simulation temperatures from 1000 K to 1800 K. Each simulation contains 1000 ion pairs and is simulated for 100 ps. Interactions are modeled by a Born-Meyer-Huggins-Tosi-Fumi (BMHTF) potential (Pan et al., 2016; Fumi and Tosi, 1964; Huggins and Mayer, 1933) and are analysed using ASE (Larsen et al., 2017). The experiments are queued using hydra parameter composition and deployed in parallel using dask4dvc. Using the DVC extension for Visual Studio Code, we can follow the progress of the experiments, e.g. the temperature, in real time.

⬇

1dvc exp run --queue -S "lammps.yaml:temperature=range(1000,1800,50)"

2dask4dvc run --config config.yaml

Figure 6. Parallel execution of multiple experiments.

After an initial structure has been generated, the respective nodes that generate the box are cached and the MD simulations are executed in parallel. Figure 6 visualizes the parts of the experiments that are only executed once and the parts that are executed for each experiment with different parameters.

For further comparison, the results of each experiment can be visualized using the DVC plot feature. A more detailed analysis can be performed using the ZnTrack package. Here, we can easily iterate over the experiments and load the results without making changes to the workspace. Figure 7 shows the density of the systems as well as the radial distribution function (RDF) $g(r)$, in which temperature dependent structural changes can be seen.

Figure 7. Density and RDF of molten sodium chloride at different temperatures. The RDFs are displayed for the Na-Cl pair.

Every experiment is referenced by a unique identifier and is associated with a commit. All experiments are made available through the GIT repository and can be listed using dvc exp list <remote>. The ZnTrack package allows us to load the experiments directly. The only requirement for this to work is that the respective package for the Lammps node is installed. This node can then be inspected or connected to another computational graph.

⬇

1node = zntrack.from_rev("Lammps", rev="exp1")

2print(node.thermo)

5.1.2. Structural Database Repository

We have demonstrated how to use ZnTrack to run an atomistic simulation using Lammps and a BMHTF potential. In recent years, machine learning has become a popular tool to model interatomic interactions instead. These models are often trained on simulations from first principle methods. Often, these datasets are shared in a way that either does not allow versioning them or will inflate the size of the repository by continuously editing large binary files that are not designed to be stored via GIT (Eastman et al., 2023; Xiao et al., 2017). Many of these datasets are designed as benchmarks for machine learning models and therefore should not be altered. They are also used to train production models (Ramakrishnan et al., 2014; Chmiela et al., 2017; Smith et al., 2017).

With the structural database repository (SDR) we want to introduce a GIT repository with a public DVC remote that can be used to access, comment on and expand or improve first principle simulation data. All information required to reproduce the data is available through this repository. Data is stored in the H5MD format (Buyl et al., 2014) and can be accessed either directly or through ZnTrack using ASE (Larsen et al., 2017).

⬇

1import zntrack

2

3node = zntrack.from_rev(

4    "BMIM_BF4_x10",

5    remote="https://github.com/user/repo",

6    rev="water"

7)

8node.atoms: list[ase.Atoms]

Furthermore, tools can be built directly on top of ZnTrack to make working with the data even easier. We provide the ZnDraw package that can visualize not only most molecular dynamics trajectory files but also access the SDR directly as shown in Figure 8.

⬇

1zndraw BMIM_BF4_x10-x10.atoms --remote https://github.com/user/repo --rev BMIM-X

Figure 8. The room temperature ionic liquid BMIM-BF${}_{\textrm{4}}$ visualized from the SDR using ZnDraw.

5.2. Random Forest Classifier

ZnTrack can also be used for more general machine learning applications. We will showcase this by implementing a machine learning method for a binary classification task. To this end, we will demonstrate the training of a random forest classifier (Breiman, 2001) adapted from the DVC example repository (Iterative, 2023c) in ZnTrack. This example utilizes a dataset consisting of StackOverflow questions and their respective tags.

Figure 9. Computational graph for training a random forest classifier.

The code is organized into distinct steps, as illustrated in Figure 9. The initial step involves extracting data from the raw data file and converting it into a suitable format. Following this, the data undergoes featurization using a bag-of-words approach, and it is subsequently divided into separate training and test sets. The features extracted from the training set are then employed to train a random forest classifier. Lastly, the model’s performance is evaluated using the unseen test set. The code is implemented in Python and uses scikit-learn (Pedregosa et al., 2011), numpy (Harris et al., 2020), pandas (pandas development team, 2020) and matplotlib (Hunter, 2007).

⬇

1params = yaml.safe_load(open("params.yaml"))["train"]

2

3inputs = sys.argv[1]

4output = sys.argv[2]

5seed = params["seed"]

6n_est = params["n_est"]

7min_split = params["min_split"]

We’ve converted the example code (see above) to adhere to the DaC paradigm by using ZnTrack (see code below). For this scenario, we were able to reduce the lines of code and the McCabe complexity (McCabe, 1976) between parts of the workflow with the same functionality. Adapting data into the DaC paradigm through ZnTrack often only requires defining parameters, dependencies and outputs as class attributes, whilst the remaining code can be kept almost unchanged.

⬇

1class Train(zntrack.Node):

2    min_split: float = zntrack.params(0.01)

3    n_est: int = zntrack.params(50)

4    seed: int = zntrack.params(20170428)

5

6    features: str = zntrack.deps_path("data/features")

7    model: str = zntrack.outs_path("model.pkl")

This is further supported by the similarity analysis shown in Figure 10. In this analysis, we compute TF-IDF features (Manning et al., 2008) for both the original code and the code written with ZnTrack. We then calculate the cosine similarity between them. The results indicate that, despite adaptations, the new code remains similar to the original. This strengthens the argument that ZnTrack can be used to convert existing code to adhere to the DaC paradigm, while keeping many of the original code’s characteristics, thereby only requiring minimal changes and being not very invasive.

Figure 10. Similarity between the original code and the adapted code for DaC using ZnTrack. The similarity is calculated using the cosine similarity between the TF-IDF features (Manning et al., 2008) as implemented in scikit-learn (Pedregosa et al., 2011).

5.3. Data Availability

The described packages and conducted experiments are publicly available under the Apache-2.0 license.

• ZnTrack https://github.com/zincware/ZnTrack and pip install zntrack.

• dask4dvc https://github.com/zincware/dask4dvc and pip install dask4dvc

• ZnDraw https://github.com/zincware/ZnDraw and pip install zndraw

• Molecular Dynamics Simulations https://github.com/PythonFZ/zntrack_lammps

• Random Forest Classifier https://github.com/PythonFZ/DataAsCodeArtifacts

6. Conclusion

In this work, we introduced the DaC concept, which simplifies access, sharing, and version control of data and workflows. DaC shifts the interface between data and code from end-users to developers and provides it alongside the code. We present the ZnTrack package, which facilitates DaC implementation in Python.

We demonstrated how DaC can be applied in atomistic simulations, allowing for easy parallel execution of multiple simulations and the FAIR sharing and version control of large datasets. Furthermore, we showed how packages can be built to interface with other DaC tools, as briefly demonstrated by the ZnTrack interface of the ZnDraw visualizer. Fully reproducible workflows support entry into data-driven fields, such as machine learning and atomistic simulations.

Leveraging DVC and GIT, ZnTrack offers easy integration of DaC into the research software and data lifecycle. It accommodates both newcomers and experienced researchers, providing minimal overhead for converting existing workflows.

As an open-source package, ZnTrack also serves as a flexible workflow framework. It enables DaC within Python, promoting the close integration of code and associated data. This strengthens FAIR data practices and fosters collaborations within and beyond research groups. Publishing research software, datasets, and experiments on platforms like GitHub or GitLab can contribute to the growth of open science, as seen in open-source research software.