Skip to content

API Reference

average_efficiency(coords_donor, coords_acceptor, r0=50.0)

Compute ensemble-averaged FRET efficiency.

Parameters:

Name Type Description Default
coords_donor Float[Array, 'M 3']

(M, 3) donor coordinates for M frames.

 required
coords_acceptor Float[Array, 'M 3']

(M, 3) acceptor coordinates for M frames.

 required
r0 float

Förster distance.

 50.0

Returns:

Type Description
Float[Array, '']

Scalar averaged efficiency .
Source code in diff_fret/kernels.py 
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
def average_efficiency(
    coords_donor: Float[Array, "M 3"],
    coords_acceptor: Float[Array, "M 3"],
    r0: float = 50.0,
) -> Float[Array, ""]:
    """
    Compute ensemble-averaged FRET efficiency.

    Args:
        coords_donor: (M, 3) donor coordinates for M frames.
        coords_acceptor: (M, 3) acceptor coordinates for M frames.
        r0: Förster distance.

    Returns:
        Scalar averaged efficiency <E>.
    """
    r = distance_distribution(coords_donor, coords_acceptor)
    e = fret_efficiency(r, r0)
    return jnp.mean(e)

distance_distribution(donor_coords, acceptor_coords)

Compute donor-acceptor distances.

Parameters:

Name Type Description Default
donor_coords Float[Array, '*batch 3']

(*batch, 3) coordinates of donor(s).

 required
acceptor_coords Float[Array, '*batch 3']

(*batch, 3) coordinates of acceptor(s).

 required

Returns:

Type Description
Float[Array, '*batch']

Distances (*batch,).
Source code in diff_fret/kernels.py 
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
def distance_distribution(
    donor_coords: Float[Array, "*batch 3"],
    acceptor_coords: Float[Array, "*batch 3"],
) -> Float[Array, "*batch"]:
    """
    Compute donor-acceptor distances.

    Args:
        donor_coords: (*batch, 3) coordinates of donor(s).
        acceptor_coords: (*batch, 3) coordinates of acceptor(s).

    Returns:
        Distances (*batch,).
    """
    dist_sq = jnp.sum((donor_coords - acceptor_coords) ** 2, axis=-1)
    # Safe distance for gradients (avoids NaN at dist=0)
    dist = jnp.sqrt(jnp.where(dist_sq > 0, dist_sq, 1.0))
    return jnp.where(dist_sq > 0, dist, 0.0)

fret_efficiency(r, r0=50.0)

Compute FRET efficiency using Förster theory.

Parameters:

Name Type Description Default
r Float[Array, '*batch']

Distance(s) in Angstroms.

 required
r0 float

Förster distance in Angstroms (default 50.0).

 50.0

Returns:

Type Description
Float[Array, '*batch']

FRET efficiency E in [0, 1].
Source code in diff_fret/kernels.py 
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
def fret_efficiency(
    r: Float[Array, "*batch"],
    r0: float = 50.0,
) -> Float[Array, "*batch"]:
    """
    Compute FRET efficiency using Förster theory.

    Args:
        r: Distance(s) in Angstroms.
        r0: Förster distance in Angstroms (default 50.0).

    Returns:
        FRET efficiency E in [0, 1].
    """
    return 1.0 / (1.0 + (r / r0) ** 6)

fret_efficiency_av(attachment_donor, attachment_acceptor, key, radius_donor=10.0, radius_acceptor=10.0, n_samples=50, r0=50.0)

Differentiable Accessible Volume (AV) simulation for FRET. Models the dye as a Gaussian spatial distribution around its attachment point.

Parameters:

Name Type Description Default
attachment_donor Float[Array, 3]

(3,) attachment point for donor.

 required
attachment_acceptor Float[Array, 3]

(3,) attachment point for acceptor.

 required
key Array

JAX PRNG key.

 required
radius_donor float

Characteristic radius (standard deviation) of donor dye distribution (default 10.0).

 10.0
radius_acceptor float

Characteristic radius (standard deviation) of acceptor dye distribution (default 10.0).

 10.0
n_samples int

Number of samples to use for the Monte Carlo integration (default 50).

 50
r0 float

Förster distance in Angstroms (default 50.0).

 50.0

Returns:

Type Description
Float[Array, '']

Averaged efficiency over the accessible volumes.
Source code in diff_fret/kernels.py 
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
def fret_efficiency_av(
    attachment_donor: Float[Array, "3"],
    attachment_acceptor: Float[Array, "3"],
    key: jax.Array,
    radius_donor: float = 10.0,
    radius_acceptor: float = 10.0,
    n_samples: int = 50,
    r0: float = 50.0,
) -> Float[Array, ""]:
    """
    Differentiable Accessible Volume (AV) simulation for FRET.
    Models the dye as a Gaussian spatial distribution around its attachment point.

    Args:
        attachment_donor: (3,) attachment point for donor.
        attachment_acceptor: (3,) attachment point for acceptor.
        key: JAX PRNG key.
        radius_donor: Characteristic radius (standard deviation) of donor dye distribution (default 10.0).
        radius_acceptor: Characteristic radius (standard deviation) of acceptor dye distribution (default 10.0).
        n_samples: Number of samples to use for the Monte Carlo integration (default 50).
        r0: Förster distance in Angstroms (default 50.0).

    Returns:
        Averaged efficiency <E> over the accessible volumes.
    """
    key_d, key_a = jax.random.split(key)

    # Sample dye positions from Gaussian clouds
    # We use the reparameterization trick (attachment + radius * noise) for differentiability
    noise_d = jax.random.normal(key_d, (n_samples, 3))
    noise_a = jax.random.normal(key_a, (n_samples, 3))

    pos_d = attachment_donor + radius_donor * noise_d
    pos_a = attachment_acceptor + radius_acceptor * noise_a

    # Compute paired-sample distances: (n_samples, 3) → (n_samples,)
    # The double integral ∫∫ P_D(r_D) P_A(r_A) E(|r_D - r_A|) dr_D dr_A
    # factorises for independent clouds and is correctly estimated by N paired draws,
    # NOT by the N² all-to-all cross-product (which is O(N²) and over-counts).
    diff = pos_d - pos_a  # (n_samples, 3)
    dist = jnp.sqrt(jnp.sum(diff**2, axis=-1) + 1e-9)  # (n_samples,)

    efficiencies = fret_efficiency(dist, r0)

    return jnp.mean(efficiencies)

kappa_squared_bounds(anisotropy_donor, anisotropy_acceptor)

Calculate the bounds of the orientation factor kappa^2 based on Dale-Eisinger-Blumberg (1979) theory.

Parameters:

Name Type Description Default
anisotropy_donor float

Measured steady-state anisotropy of donor.

 required
anisotropy_acceptor float

Measured steady-state anisotropy of acceptor.

 required

Returns:

Type Description
Float[Array, 2]

(min_kappa2, max_kappa2) array.
Source code in diff_fret/kernels.py 
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
def kappa_squared_bounds(
    anisotropy_donor: float,
    anisotropy_acceptor: float,
) -> Float[Array, "2"]:
    """
    Calculate the bounds of the orientation factor kappa^2 based on
    Dale-Eisinger-Blumberg (1979) theory.

    Args:
        anisotropy_donor: Measured steady-state anisotropy of donor.
        anisotropy_acceptor: Measured steady-state anisotropy of acceptor.

    Returns:
        (min_kappa2, max_kappa2) array.
    """
    # Limiting anisotropy (r0) typically 0.4 for standard dyes
    r0_limit = 0.4

    # Depolarization factors d = sqrt(r/r0)
    d_d = jnp.sqrt(jnp.clip(anisotropy_donor / r0_limit, 0.0, 1.0))
    d_a = jnp.sqrt(jnp.clip(anisotropy_acceptor / r0_limit, 0.0, 1.0))

    # Max/Min bounds for dynamic dynamic regime
    min_k2 = (2.0 / 3.0) * (1.0 - 0.5 * (d_d + d_a))
    max_k2 = (2.0 / 3.0) * (1.0 + d_d + d_a + 3.0 * d_d * d_a)

    return jnp.array([min_k2, max_k2])