Inheritance

ZnTrack allows inheritance from a Node base class. This can e.g. be useful if you want to test out different methods of the same kind. In the following example, we will show this by using different functions in the run method with the same inputs and outputs.

from zntrack import config

config.nb_name = "02_Inheritance.ipynb"

# Work in a temporary directory
from zntrack.utils import cwd_temp_dir

temp_dir = cwd_temp_dir()

!git init
!dvc init

Initialized empty Git repository in /tmp/tmp4o5jiism/.git/
Initialized DVC repository.

You can now commit the changes to git.

+---------------------------------------------------------------------+
|                                                                     |
|        DVC has enabled anonymous aggregate usage analytics.         |
|     Read the analytics documentation (and how to opt-out) here:     |
|             <https://dvc.org/doc/user-guide/analytics>              |
|                                                                     |
+---------------------------------------------------------------------+

What's next?
------------
- Check out the documentation: <https://dvc.org/doc>
- Get help and share ideas: <https://dvc.org/chat>
- Star us on GitHub: <https://github.com/iterative/dvc>

import zntrack

Let us define a NodeBase which has a single input and output. We will use this as a base class for the following Nodes.

• AddNumber: Shift input by an offset

• MultiplyNumber: Multiply input by a factor

Both of these Nodes extend the NodeBase by additional parameters.

class NodeBase(zntrack.Node):
    _name_ = "basic_number"

    inputs: float = zntrack.zn.params()
    output: float = zntrack.zn.outs()

class AddNumber(NodeBase):
    """Shift input by an offset"""

    offset: float = zntrack.zn.params()

    def run(self):
        self.output = self.inputs + self.offset


class MultiplyNumber(NodeBase):
    """Multiply input by a factor"""

    factor: float = zntrack.zn.params()

    def run(self):
        self.output = self.inputs * self.factor

with zntrack.Project() as project:
    add_number = AddNumber(inputs=10.0, offset=15.0)
project.run()

Running DVC command: 'stage add --name basic_number --force ...'

Creating 'dvc.yaml'
Adding stage 'basic_number' in 'dvc.yaml'

To track the changes with git, run:

        git add dvc.yaml nodes/basic_number/.gitignore

To enable auto staging, run:

        dvc config core.autostage true

Jupyter support is an experimental feature! Please save your notebook before running this command!
Submit issues to https://github.com/zincware/ZnTrack.
[NbConvertApp] Converting notebook 02_Inheritance.ipynb to script

Running stage 'basic_number':
> zntrack run src.AddNumber.AddNumber --name basic_number

[NbConvertApp] Writing 4756 bytes to 02_Inheritance.py

Generating lock file 'dvc.lock'
Updating lock file 'dvc.lock'

To track the changes with git, run:

        git add dvc.lock

To enable auto staging, run:

        dvc config core.autostage true
Use `dvc push` to send your updates to remote storage.

Because the Nodes inherit from each other and we defined the node_name in the parent class, we can use all classes to load the outputs (as long as they are shared). This is important to keep in mind when working with inheritance, that the output might not necessarily be created by the Node it was loaded by. On the other hand, this can be handy for dependency handling. A subsequent Node can e.g. depend on the parent Node and does not need to know where the values actually come from. I.e. an ML Model might implement a predict function in the parent node but can have an entirely different structure. An evaluation node might only need the predict method and can therefore be used with all children of the model class.

NodeBase.from_rev().output

25.0

!dvc dag

+--------------+
| basic_number |
+--------------+

with zntrack.Project() as project:
    multiply_number = MultiplyNumber(inputs=6.0, factor=6.0)
project.run()

Running DVC command: 'stage add --name basic_number --force ...'

Modifying stage 'basic_number' in 'dvc.yaml'

To track the changes with git, run:

        git add dvc.yaml

To enable auto staging, run:

        dvc config core.autostage true

[NbConvertApp] Converting notebook 02_Inheritance.ipynb to script

Running stage 'basic_number':
> zntrack run src.MultiplyNumber.MultiplyNumber --name basic_number

[NbConvertApp] Writing 4756 bytes to 02_Inheritance.py

Updating lock file 'dvc.lock'

To track the changes with git, run:

        git add dvc.lock

To enable auto staging, run:

        dvc config core.autostage true
Use `dvc push` to send your updates to remote storage.

NodeBase.from_rev().output

36.0

As expected the node name remains the same and therefore, the Node is replaced with the new one.

!dvc dag

+--------------+
| basic_number |
+--------------+

Nodes as parameters

Sometimes it can be useful to have a Node as a parameter or use the run method of the given Node but storing the outputs somewhere else. For example an active learning cycle might use the model and evaluation class but the outputs are stored in the active learning Node. You might still want to use the other Nodes to avoid overhead though.

In the following we will use the run method of a NodeBase Node and also have a dataclass Node just for storing parameters. Internally, ZnTrack disables all outputs of the given Node except for a UUID file.

class DivideNumber(NodeBase):
    """Multiply input by a factor"""

    divider: float = zntrack.zn.params()

    def run(self):
        self.output = self.inputs * self.divider


class Polynomial(zntrack.Node):
    a0: float = zntrack.zn.params()
    a1: float = zntrack.zn.params()


class ManipulateNumber(zntrack.Node):
    inputs: float = zntrack.zn.params()
    output: float = zntrack.zn.outs()
    value_handler: NodeBase = zntrack.zn.nodes()
    polynomial: Polynomial = zntrack.zn.nodes()

    def run(self):
        # use the passed method
        self.value_handler.inputs = self.inputs
        self.value_handler.run()
        self.output = self.value_handler.output
        # polynomials
        self.output = self.polynomial.a0 + self.polynomial.a1 * self.output

value_handler = DivideNumber(divider=3.0, inputs=None)
polynomial = Polynomial(a0=60.0, a1=10.0)

with zntrack.Project() as project:
    manipulate_number = ManipulateNumber(
        inputs=10.0,
        value_handler=value_handler,
        polynomial=polynomial,
    )
project.run()

Running DVC command: 'stage add --name ManipulateNumber --force ...'

Adding stage 'ManipulateNumber' in 'dvc.yaml'

To track the changes with git, run:

        git add nodes/ManipulateNumber/.gitignore dvc.yaml

To enable auto staging, run:

        dvc config core.autostage true

Running DVC command: 'stage add --name ManipulateNumber_polynomial --outs ...'

Adding stage 'ManipulateNumber_polynomial' in 'dvc.yaml'

Could not create .gitignore entry in /tmp/tmp4o5jiism/nodes/ManipulateNumber_polynomial/.gitignore. DVC will attempt to create .gitignore entry again when the stage is run.

To track the changes with git, run:

        git add dvc.yaml

To enable auto staging, run:

        dvc config core.autostage true

Running DVC command: 'stage add --name ManipulateNumber_value_handler --outs ...'

Adding stage 'ManipulateNumber_value_handler' in 'dvc.yaml'

Could not create .gitignore entry in /tmp/tmp4o5jiism/nodes/ManipulateNumber_value_handler/.gitignore. DVC will attempt to create .gitignore entry again when the stage is run.

To track the changes with git, run:

        git add dvc.yaml

To enable auto staging, run:

        dvc config core.autostage true

[NbConvertApp] Converting notebook 02_Inheritance.ipynb to script
[NbConvertApp] Writing 4756 bytes to 02_Inheritance.py
[NbConvertApp] Converting notebook 02_Inheritance.ipynb to script
[NbConvertApp] Writing 4756 bytes to 02_Inheritance.py
[NbConvertApp] Converting notebook 02_Inheritance.ipynb to script

Running stage 'ManipulateNumber_polynomial':
> zntrack run src.Polynomial.Polynomial --name ManipulateNumber_polynomial --hash-only

[NbConvertApp] Writing 4756 bytes to 02_Inheritance.py

Updating lock file 'dvc.lock'

Running stage 'ManipulateNumber_value_handler':
> zntrack run src.DivideNumber.DivideNumber --name ManipulateNumber_value_handler --hash-only
Updating lock file 'dvc.lock'

Running stage 'ManipulateNumber':
> zntrack run src.ManipulateNumber.ManipulateNumber --name ManipulateNumber
Updating lock file 'dvc.lock'

Stage 'basic_number' didn't change, skipping

To track the changes with git, run:

        git add nodes/ManipulateNumber_value_handler/.gitignore nodes/ManipulateNumber_polynomial/.gitignore dvc.lock

To enable auto staging, run:

        dvc config core.autostage true
Use `dvc push` to send your updates to remote storage.

manipulate_number.load()

manipulate_number.output

360.0

!dvc dag

+--------------+
| basic_number |
+--------------+
+-----------------------------+                    +--------------------------------+
| ManipulateNumber_polynomial |                    | ManipulateNumber_value_handler |
+-----------------------------+                    +--------------------------------+
                           ****                      *****
                               ****              ****
                                   ***        ***
                                +------------------+
                                | ManipulateNumber |
                                +------------------+

temp_dir.cleanup()