Extending Emission Clauses

This notebook shows how to create new fluent emission methods by subclassing nkdsl.AbstractEmissionClause.

Use this extension point when you want reusable branch templates (for example, “emit with this guard and this standard tag/update policy”) so you do not repeat boilerplate emit_if(...) blocks everywhere.

Emission clause contract

A custom emission clause should return nkdsl.EmissionClauseSpec. The key fields are:

• mode: one of emit, emit_if, emit_elseif, emit_else

• predicate: branch predicate (used by conditional modes)

• update: rewrite program (identity, shift, write, chained updates, …)

• matrix_element: amplitude expression for that branch

• tag: optional diagnostic tag

In practice, most custom clauses are thin wrappers that package a repeated policy into one method.

import jax.numpy as jnp
import netket as nk
import nkdsl

hi = nk.hilbert.Fock(n_max=3, N=4)
print("Hilbert size:", hi.size)

∣NK⟩ Tip: If timeit=True signals high \% spent sampling n_discarded, consider lowering it.

Hilbert size: 4

Example 1: emit_when_at_least(...)

This clause emits only when a label’s occupancy is at least a cutoff. It returns an emit_if spec so it can be chained with built-in emit_elseif or emit_else.

class EmitWhenAtLeast(nkdsl.AbstractEmissionClause):
    clause_name = "emit_when_at_least"

    def build_emission(self, ctx, label: str = "i", cutoff: int = 1):
        predicate = ctx.site(label).value >= int(cutoff)
        return nkdsl.EmissionClauseSpec(
            mode="emit_if",
            predicate=predicate,
            update=nkdsl.identity(),
            matrix_element=ctx.site(label).value,
            tag="emit-when-at-least",
        )


nkdsl.register_emission_clause(EmitWhenAtLeast, replace=True)
print("Registered:", "emit_when_at_least" in nkdsl.available_emission_clause_names())

Registered: True

op_custom_sym = (
    nkdsl.SymbolicDiscreteJaxOperator(hi, "custom-emission")
    .for_each_site("i")
    .emit_when_at_least("i", cutoff=2)
    .emit_else(nkdsl.identity(), matrix_element=0.0, tag="fallback")
    .build()
)
op_custom = op_custom_sym.compile()

_xp, mels = op_custom.get_conn_padded(jnp.asarray([0, 1, 2, 3], dtype=jnp.int32))
print("mels:", mels)

mels: [5. 0. 0. 0. 0. 0. 0. 0.]

print(op_custom_sym.to_ir())

symbolic.operator @"custom-emission" [dtype=float64, hermitian=false, hilbert_size=4] {
  ; 1 term(s)

  term #0 "0" [kbody, n_iter=4, max_conn_size=8] {
    iterate: for (i,) in [(0,), (1,), (2,), ... +1 more]
    where:   true
    emit #0 [tag='emit-when-at-least']:
      where:     (x[i] >= 2)
      update:    identity
      amplitude: x[i]
    emit #1 [tag='fallback']:
      where:     !(x[i] >= 2)
      update:    identity
      amplitude: 0
  }

}

Notice how one custom fluent method captures both a predicate policy and default branch metadata. This is the main value proposition: you call one readable method, while internal behavior stays precise.

Example 2: registration via generic @register

The generic decorator can register iterator, predicate, or emission clauses. Here we use it for a second emission clause to show a compact authoring path.

@nkdsl.register
class EmitIfNonZero(nkdsl.AbstractEmissionClause):
    clause_name = "emit_if_nonzero"

    def build_emission(self, ctx, label: str = "i", *, mel: float = 1.0):
        return nkdsl.EmissionClauseSpec(
            mode="emit_if",
            predicate=ctx.site(label).value != 0,
            update=nkdsl.identity(),
            matrix_element=float(mel),
            tag="nonzero",
        )


print("Registered:", "emit_if_nonzero" in nkdsl.available_emission_clause_names())

op_nonzero_sym = (
    nkdsl.SymbolicDiscreteJaxOperator(hi, "nonzero")
    .for_each_site("i")
    .emit_if_nonzero("i", mel=5.0)
    .emit_else(nkdsl.identity(), matrix_element=0.0)
    .build()
)
op_nonzero = op_nonzero_sym.compile()

_xp2, m2 = op_nonzero.get_conn_padded(jnp.asarray([0, 1, 0, 2], dtype=jnp.int32))
print("m2:", m2)

Registered: True
m2: [20.  0.  0.  0.  0.  0.  0.  0.]

Design rules for production emission clauses

• Keep one clause focused on one branch policy.

• Validate arguments inside build_emission.

• Choose mode deliberately and document it.

• Include tags when they improve diagnostics.

• Print IR during development to verify emitted branch structure.

A good emission clause should make common logic shorter, clearer, and harder to misuse.

Where this fits in the extension architecture

With custom iterators, predicates, and emissions together, you can shape:

• where to iterate (iterator clause)

• when to activate (predicate clause)

• what to emit (emission clause)

That separation keeps DSL extensions composable and lets teams standardize domain-specific patterns without forking the compiler.