Objectives & Loss Functions
Build custom loss functions using composable expression trees. From single-mode coupling to broadband multi-port objectives with penalty terms.
What you will learn
- 1.Build objectives from measurement primitives (mode coupling, power, field)
- 2.Compose broadband and multi-port loss functions with operator overloading
- 3.Add penalty terms for reflection suppression and extinction ratio
- 4.Understand serialization: how expression trees become JSON for GPU evaluation
The Default Objective
In the inverse design tutorial, you optimized a grating coupler using a single line: hwc.objectives.mode_coupling(...). That function returned an expression tree node, a pure-data description of what to compute, not the computation itself. In this tutorial, you will learn to compose your own objectives from scratch, building everything from simple single-mode coupling up to multi-objective loss functions with penalty terms.
Prerequisites: you should have completed the Inverse Design tutorial and have a working API key. Run the first cell below to set up the shared variables (mode_field, input_power, mode_cross_power) used by all subsequent cells. You do not need a running optimization to follow along.
import numpy as np
import hyperwave_community as hwc
from hyperwave_community import objectives as obj
# Solve the TE0 mode (same as inverse design tutorial)
mode_field, n_eff = hwc.solve_waveguide_mode(
grid=0.035, waveguide_width=0.5, waveguide_height=0.220,
n_core=3.48, n_clad=1.44, wavelength=1.55, mode_number=0,
)
input_power = 1.0
mode_cross_power = 1.0
# The default objective: mode coupling efficiency
eff = obj.mode_coupling(
mode_field=mode_field,
input_power=input_power,
mode_cross_power=mode_cross_power,
monitor="waveguide_output",
)
print(f"Type: {type(eff).__name__}")
print(f"Repr: {repr(eff)}")The return value is a ModeCoupling node, not a number. Nothing has been computed yet. The tree describes *what* to measure; the GPU evaluates it during optimization.
Why Expression Trees?
You might wonder: why not just pass a Python function like lambda fields: jnp.sum(jnp.abs(fields['wg'][0,1])**2)? Two reasons.
Hyperwave runs optimizations on shared GPU infrastructure. Accepting arbitrary Python functions would mean running user code on the GPU, a security risk in a multi-tenant environment. Expression trees are pure data: a fixed set of operations that the server's JAX interpreter can evaluate safely.
Expression trees serialize to JSON. Your objective can be saved, restored, inspected, and sent over the network. A Python lambda cannot.
The trade-off is that you can only use the operations provided by the objectives module. But as you will see, the available primitives cover the vast majority of photonics optimization problems.
# Expression trees compose with Python operators.
# Behind the scenes, __add__, __mul__, etc. build tree nodes.
a = obj.mode_coupling(mode_field, input_power, mode_cross_power, "port1")
b = obj.mode_coupling(mode_field, input_power, mode_cross_power, "port2")
combined = a + b # Add node
weighted = 0.7 * a # Mul(Const(0.7), a)
negated = -a # Neg node
scaled = a ** 2 # Pow node
print(f"a + b: {combined}")
print(f"0.7 * a: {weighted}")
print(f"-a: {negated}")
print(f"a ** 2: {scaled}")Measurement Primitives
Every objective starts from measurement nodes, leaf nodes that read data from simulation monitors. There are four measurement types, from high-level to low-level:
Bidirectional overlap integral with a reference mode. Returns a scalar between 0 and 1. This is what you used in the inverse design tutorial.
Poynting vector power flow through a monitor plane. Returns total integrated power along a given axis. Useful for total transmission or reflection measurements.
Integrated |E|^2 for a single E-field component (Ex, Ey, or Ez) over a monitor. Useful for focusing objectives where you want to maximize energy density.
Raw complex field component at a monitor. Returns the full spatial array, not a scalar. This is the most flexible primitive. Combine with sum_spatial, real, conj, etc. to build any custom measurement.
Start with mode_coupling or power when possible. They are physically meaningful and numerically stable. Drop down to field only when you need spatial control over the objective.
# All four measurement types
m1 = obj.mode_coupling(mode_field, input_power, mode_cross_power, "wg_out")
m2 = obj.power("transmission_plane", axis=0)
m3 = obj.intensity("Ey", "focus_monitor")
m4 = obj.field("Ey", "focus_monitor")
# field() returns a spatial array -- must reduce to scalar for optimization
focusing = obj.sum_spatial(obj.abs_val(m4) ** 2)
print(f"mode_coupling: {m1}")
print(f"power: {m2}")
print(f"intensity: {m3}")
print(f"field -> sum: {focusing}")Broadband Optimization (WDM Demux)
A single-wavelength objective is fragile: the optimizer can exploit narrowband resonances that fall apart at adjacent wavelengths. For broadband devices like WDM demultiplexers, you want to maximize the worst-case coupling across multiple wavelengths.
min_of() takes the minimum across its arguments. Maximizing this minimum forces the optimizer to bring up the weakest channel, producing flat broadband performance instead of a single sharp peak.
# WDM demux: 4 channels, each coupling to a different output port
# Wavelengths: 1530, 1540, 1550, 1560 nm (20nm spacing, CWDM-like)
wavelengths = [1.530, 1.540, 1.550, 1.560]
port_names = ["drop1", "drop2", "drop3", "drop4"]
# One mode_coupling per channel, each at a different port and frequency
channels = []
for i, (wl, port) in enumerate(zip(wavelengths, port_names)):
ch = obj.mode_coupling(
mode_field=mode_field,
input_power=input_power,
mode_cross_power=mode_cross_power,
monitor=port,
freq_idx=i, # each wavelength is a separate frequency index
)
channels.append(ch)
# Maximize the worst-performing channel
objective = obj.min_of(*channels)
print(f"Broadband objective: {objective}")
print(f"Number of channels: {len(channels)}")The multi-frequency simulation is set up via freq_band in hwc.build_device(). Each freq_idx corresponds to one frequency in that band. The optimizer runs a single FDTD simulation that captures all frequencies simultaneously (broadband source), then evaluates the objective at each frequency index.
Multi-Port Balancing (Power Splitter)
A 1x2 power splitter should deliver 50% of the input power to each output port. A naive objective like coupling_port1 + coupling_port2 allows the optimizer to dump all power into one port (100% + 0% = 100%). Instead, we want to penalize imbalance.
Two strategies work well: (a) maximize the minimum of the two port efficiencies, or (b) add the two efficiencies and subtract a penalty for their difference.
# 1x2 power splitter: two output waveguides
port1 = obj.mode_coupling(mode_field, input_power, mode_cross_power, "out_top")
port2 = obj.mode_coupling(mode_field, input_power, mode_cross_power, "out_bot")
# Strategy A: maximize worst-case port
objective_a = obj.min_of(port1, port2)
# Strategy B: maximize total with balance penalty
total = port1 + port2
imbalance = obj.abs_val(port1 - port2)
objective_b = total - 2.0 * imbalance
print(f"Strategy A (min): {objective_a}")
print(f"Strategy B (penalize): {objective_b}")Strategy A is simpler and usually sufficient. Strategy B gives you a tunable knob: the weight on the imbalance penalty (2.0 here) controls how aggressively the optimizer enforces balance. Higher weight means tighter balance but potentially lower total throughput.
Field-Level Objectives (Metalens Focusing)
For a metalens, the objective is to concentrate the Poynting flux at a focal point. This requires working with raw field components rather than mode overlaps. The field() primitive gives you access to individual E and H components, and you compose them with complex math to build a Poynting flux integral.
When using field(), the result is a spatial array (same shape as the monitor plane). You must reduce it to a scalar using sum_spatial() or mean_spatial() before the optimizer can use it. Forgetting the reduction is a common mistake.
# Metalens focusing: maximize Poynting flux through a small focal monitor
# The focal monitor is a small plane centered at the desired focus.
# Poynting flux S_x = Re(Ey * Hz* - Ez * Hy*)
ey = obj.field("Ey", monitor="focal_plane")
ez = obj.field("Ez", monitor="focal_plane")
hy = obj.field("Hy", monitor="focal_plane")
hz = obj.field("Hz", monitor="focal_plane")
# Cross product for x-directed power flow
sx = obj.real(ey * obj.conj(hz) - ez * obj.conj(hy))
# Integrate over the focal plane
focal_power = obj.sum_spatial(sx)
print(f"Focal power objective: {focal_power}")
# Alternative: simpler approach using the power() primitive
# (equivalent if the monitor is placed correctly)
focal_power_simple = obj.power("focal_plane", axis=0)
print(f"Equivalent shorthand: {focal_power_simple}")The manual Poynting flux calculation and power() are equivalent. power() computes exactly this cross product internally. Use field() only when you need something power() cannot express: for example, maximizing intensity at a single spatial point within a larger monitor, or weighting different regions of the focal plane differently.
Penalty Terms
Real devices have constraints beyond raw efficiency. You may want to suppress back-reflections, enforce a minimum extinction ratio, or penalize unwanted radiation. Penalty terms let you add soft constraints to the objective using relu(), the differentiable max(0, x) function.
relu(threshold - measurement) is zero when the measurement exceeds the threshold, and grows linearly when it falls below. This pushes the optimizer to satisfy the constraint without making it a hard boundary that could cause convergence issues.
# Grating coupler with reflection suppression
forward = obj.mode_coupling(mode_field, input_power, mode_cross_power, "wg_output")
reflected_power = obj.power("reflection_monitor", axis=0)
# Penalize if reflection exceeds 5% of input power (-13 dB)
reflection_penalty = obj.relu(reflected_power - 0.05)
# Combined objective: maximize coupling, suppress reflection
# Weight of 5.0 on penalty makes reflection suppression aggressive
objective = forward - 5.0 * reflection_penalty
print(f"Objective with penalty: {objective}")
# Another example: extinction ratio for a switch
on_state = obj.mode_coupling(mode_field, input_power, mode_cross_power, "output")
off_state = obj.mode_coupling(mode_field, input_power, mode_cross_power, "output", freq_idx=1)
# Penalize if extinction ratio < 20 dB (on/off ratio < 100)
er_penalty = obj.relu(off_state - on_state / 100.0)
switch_objective = on_state - 10.0 * er_penalty
print(f"Switch objective: {switch_objective}")Penalty weights require tuning. Too low and the constraint is ignored; too high and the optimizer focuses entirely on satisfying the constraint at the expense of the primary objective. Start with a weight of 1.0 and increase if the constraint is not met.
Putting It All Together
Here is a realistic multi-objective for a broadband grating coupler: maximize worst-case coupling across 3 wavelengths, with a reflection penalty. This is the kind of objective you would use for a production device.
# Broadband grating coupler with reflection suppression
# 3 wavelengths: 1530nm, 1550nm, 1570nm (40nm bandwidth)
# Mode coupling at each wavelength
effs = [
obj.mode_coupling(mode_field, input_power, mode_cross_power,
monitor="wg_output", freq_idx=i)
for i in range(3)
]
# Worst-case across wavelengths (broadband)
worst_case = obj.min_of(*effs)
# Reflection penalty at each wavelength
ref_penalties = [
obj.relu(obj.power("reflection", axis=0, freq_idx=i) - 0.05)
for i in range(3)
]
total_ref_penalty = ref_penalties[0] + ref_penalties[1] + ref_penalties[2]
# Final objective: maximize worst-case coupling, suppress reflection
objective = worst_case - 3.0 * total_ref_penalty
# Pass to optimize:
# results = hwc.optimize(
# device=device,
# source=source_field,
# mode=mode_field,
# objective=objective, # <-- your custom objective tree
# phase="freeform",
# n_steps=100,
# )
print(f"Objective type: {type(objective).__name__}")
print(f"Repr: {objective}")You can pass any Objective node to hwc.optimize() via the objective= parameter. When you omit it, the default is a single mode_coupling node, which is what the inverse design tutorial used.
Serialization: What the GPU Sees
When you call hwc.optimize(), your objective tree is serialized to JSON and sent to the GPU. The server deserializes it and evaluates it with JAX, computing both the loss value and its gradient via autodiff. You can inspect the serialized form with objective.serialize().
The serialized form is a nested dict of {type, ...} nodes, with numpy arrays replaced by integer indices into a separate array list. This separation keeps the JSON small and avoids encoding large arrays as text.
import json
# Simple objective for clear serialization output
eff = obj.mode_coupling(mode_field, input_power, mode_cross_power, "wg_output")
penalty = obj.relu(obj.power("reflection", axis=0) - 0.05)
objective = eff - 2.0 * penalty
# Serialize
spec, arrays = objective.serialize()
# The spec is a JSON-safe nested dict
print("Spec (JSON):")
print(json.dumps(spec, indent=2))
print(f"\nArrays: {len(arrays)} numpy array(s)")
print(f" [0] shape={arrays[0].shape}, dtype={arrays[0].dtype}")The mode field array is the only large data in the serialization. All other nodes are lightweight JSON. The server reconstructs the full expression tree from the spec and evaluates it with JAX. Every node in the tree is differentiable, so jax.grad flows through the entire objective automatically.
That is the complete objectives API. You can express any scalar function of simulation fields as a tree of measurements, math, and reductions. The key constraint is that every path from leaf to root must reduce to a scalar: the optimizer needs a single number to minimize.