Upload 8 files

modules/ai_link_flow_emulator.py

@@ -0,0 +1,162 @@
+# modules/ai_link_flow_emulator.py
+# TripAI – AI Emulator for Link Flows
+
+from __future__ import annotations
+from dataclasses import dataclass
+from typing import Any, Dict, Tuple
+
+import numpy as np
+import pandas as pd
+from sklearn.ensemble import RandomForestRegressor
+
+from .route_assignment import aon_assignment, Network
+
+
+@dataclass
+class LinkFlowEmulator:
+    """
+    AI emulator for link flows under scaled OD demand.
+
+    Attributes
+    ----------
+    model : RandomForestRegressor
+        Multi-output regressor mapping demand scale -> link flows.
+    link_ids : np.ndarray
+        IDs of links in the same order as training outputs.
+    base_total_demand : float
+        Total baseline car OD (for reference).
+    """
+    model: RandomForestRegressor
+    link_ids: np.ndarray
+    base_total_demand: float
+
+
+def _generate_training_scenarios(
+    base_od: np.ndarray,
+    network: Network,
+    n_scenarios: int = 20,
+    low_scale: float = 0.7,
+    high_scale: float = 1.3,
+) -> Tuple[np.ndarray, np.ndarray]:
+    """
+    Generate training scenarios by scaling baseline OD and performing
+    AON assignment to obtain link flows.
+
+    Parameters
+    ----------
+    base_od : np.ndarray
+        Baseline OD matrix (veh/h).
+    network : Network
+    n_scenarios : int
+        Number of random scaling scenarios.
+    low_scale : float
+        Minimum demand scale factor.
+    high_scale : float
+        Maximum demand scale factor.
+
+    Returns
+    -------
+    X : np.ndarray
+        Feature matrix of shape (n_scenarios, 1) – the demand scale.
+    Y : np.ndarray
+        Target matrix of shape (n_scenarios, n_links) – link flows.
+    """
+    n_zones = base_od.shape[0]
+    n_links = len(network.links)
+    scales = np.random.uniform(low_scale, high_scale, size=n_scenarios)
+
+    X = scales.reshape(-1, 1)
+    Y = np.zeros((n_scenarios, n_links), dtype=float)
+
+    # Build index -> (from_zone, to_zone) map to reuse AON logic
+    # We will call the existing aon_assignment with scaled OD each time.
+    # Convert base OD to DataFrame with synthetic zone index 0..n-1.
+    zones = np.arange(n_zones)
+    base_od_df = pd.DataFrame(base_od, index=zones, columns=zones)
+
+    for idx, s in enumerate(scales):
+        od_scaled = base_od_df * s
+        flows_df = aon_assignment(od_scaled, network)
+        Y[idx, :] = flows_df["flow_vehph"].to_numpy()
+
+    return X, Y
+
+
+def train_link_flow_emulator(
+    base_car_od: np.ndarray,
+    network: Network,
+    n_scenarios: int = 20,
+) -> tuple[LinkFlowEmulator, pd.DataFrame]:
+    """
+    Train a simple RandomForest-based emulator that maps a single
+    scalar 'demand scale' to resulting link flows.
+
+    Parameters
+    ----------
+    base_car_od : np.ndarray
+        Baseline car OD matrix (veh/h).
+    network : Network
+    n_scenarios : int
+        Number of training scenarios.
+
+    Returns
+    -------
+    emulator : LinkFlowEmulator
+    training_history : pd.DataFrame
+        Scenario scales and corresponding total flows, for diagnostics.
+    """
+    X, Y = _generate_training_scenarios(base_car_od, network, n_scenarios=n_scenarios)
+
+    model = RandomForestRegressor(
+        n_estimators=200,
+        max_depth=12,
+        random_state=42,
+    )
+    model.fit(X, Y)
+
+    link_ids = network.links.index.to_numpy()
+    base_total = float(base_car_od.sum())
+
+    # Build training history for inspection
+    total_flows = Y.sum(axis=1)
+    training_history = pd.DataFrame(
+        {
+            "scale": X.flatten(),
+            "total_link_flow": total_flows,
+        }
+    )
+
+    emulator = LinkFlowEmulator(
+        model=model,
+        link_ids=link_ids,
+        base_total_demand=base_total,
+    )
+    return emulator, training_history
+
+
+def predict_link_flows(
+    emulator: LinkFlowEmulator,
+    scale: float,
+    network: Network,
+) -> pd.DataFrame:
+    """
+    Predict link flows for a new demand scale using the trained emulator.
+
+    Parameters
+    ----------
+    emulator : LinkFlowEmulator
+    scale : float
+        Multiplicative scaling factor relative to baseline OD.
+    network : Network
+
+    Returns
+    -------
+    pd.DataFrame
+        Link table with predicted flows in column 'flow_vehph_emulated'.
+    """
+    X_new = np.array([[scale]], dtype=float)
+    y_pred = emulator.model.predict(X_new).flatten()
+
+    links = network.links.copy()
+    links["flow_vehph_emulated"] = y_pred
+    return links

modules/gravity_model.py

@@ -0,0 +1,108 @@
+# modules/gravity_model.py
+# TripAI – Gravity Model + OD Matrix Builder
+
+from __future__ import annotations
+import numpy as np
+import pandas as pd
+
+from .utils import iterative_proportional_fitting
+
+PURPOSES = ["HBW", "HBE", "HBS"]
+
+
+def gravity_model_doubly_constrained(
+    productions: pd.Series,
+    attractions: pd.Series,
+    travel_time: pd.DataFrame,
+    beta: float = -0.1,
+    max_iter: int = 50,
+    tol: float = 1e-6,
+) -> pd.DataFrame:
+    """
+    Doubly-constrained gravity model with IPF.
+
+    T_ij ∝ P_i * A_j * f(c_ij),
+    where f(c_ij) = exp(beta * c_ij), beta < 0
+
+    IPF is used to ensure row sums match productions and column sums
+    match attractions.
+
+    Parameters
+    ----------
+    productions : pd.Series
+        Trip productions P_i by origin TAZ.
+    attractions : pd.Series
+        Trip attractions A_j by destination TAZ.
+    travel_time : pd.DataFrame
+        Impedance matrix c_ij (minutes).
+    beta : float
+        Distance-decay parameter (negative).
+    max_iter : int
+        Maximum IPF iterations.
+    tol : float
+        Tolerance for marginal convergence.
+
+    Returns
+    -------
+    pd.DataFrame
+        OD matrix T_ij (index=origins, columns=destinations).
+    """
+    idx = productions.index
+    P = productions.values.astype(float)
+    A = attractions.values.astype(float)
+
+    c = travel_time.loc[idx, idx].values.astype(float)
+    # Impedance function
+    F = np.exp(beta * c)
+
+    # Initial gravity estimate
+    T0 = np.outer(P, A) * F
+    # Avoid all-zero rows/cols
+    T0[T0 < 0] = 0.0
+
+    T_adj = iterative_proportional_fitting(T0, P, A, max_iter=max_iter, tol=tol)
+
+    return pd.DataFrame(T_adj, index=idx, columns=idx)
+
+
+def build_all_od_matrices(
+    productions_df: pd.DataFrame,
+    attractions_df: pd.DataFrame,
+    travel_time: pd.DataFrame,
+    beta: float = -0.1,
+    max_iter: int = 50,
+    tol: float = 1e-6,
+) -> dict[str, pd.DataFrame]:
+    """
+    Build OD matrices for all purposes using the doubly-constrained gravity model.
+
+    Parameters
+    ----------
+    productions_df : pd.DataFrame
+        Columns = purposes (HBW, HBE, HBS), index = TAZ.
+    attractions_df : pd.DataFrame
+        Columns = purposes (HBW, HBE, HBS), index = TAZ.
+    travel_time : pd.DataFrame
+        Travel time matrix (minutes), index/cols = TAZ.
+    beta : float
+        Distance-decay parameter for all purposes.
+    max_iter : int
+        Maximum iterations for IPF.
+    tol : float
+        Convergence tolerance.
+
+    Returns
+    -------
+    dict[str, pd.DataFrame]
+        Mapping from purpose -> OD matrix (DataFrame).
+    """
+    od_mats: dict[str, pd.DataFrame] = {}
+
+    for purpose in PURPOSES:
+        P = productions_df[purpose]
+        A = attractions_df[purpose]
+        od_mats[purpose] = gravity_model_doubly_constrained(
+            P, A, travel_time, beta=beta, max_iter=max_iter, tol=tol
+        )
+
+    return od_mats

modules/mode_choice.py

@@ -0,0 +1,137 @@
+# modules/mode_choice.py
+# TripAI – Multinomial Logit Mode Choice
+
+from __future__ import annotations
+from dataclasses import dataclass
+from typing import Dict
+
+import numpy as np
+import pandas as pd
+
+
+@dataclass
+class ModeChoiceResult:
+    """
+    Container for mode choice outputs.
+
+    Attributes
+    ----------
+    total_od : pd.DataFrame
+        OD matrix summed over all purposes.
+    volumes : Dict[str, pd.DataFrame]
+        Mode-specific OD matrices (Car/Metro/Bus).
+    probabilities : Dict[str, pd.DataFrame]
+        Mode choice probabilities per OD pair.
+    """
+    total_od: pd.DataFrame
+    volumes: Dict[str, pd.DataFrame]
+    probabilities: Dict[str, pd.DataFrame]
+
+
+def mode_choice(
+    od_matrices: Dict[str, pd.DataFrame],
+    taz: pd.DataFrame,
+    travel_time: pd.DataFrame,
+    beta_time: float = -0.06,
+    beta_cost: float = -0.03,
+    beta_car_own: float = 0.5,
+) -> ModeChoiceResult:
+    """
+    Apply a simple Multinomial Logit (MNL) mode choice model for
+    Car / Metro / Bus.
+
+    U_m = β_time * time_m + β_cost * cost_m + γ * car_ownership (for car only)
+
+    Parameters
+    ----------
+    od_matrices : dict[str, pd.DataFrame]
+        OD matrices by purpose (HBW, HBE, HBS).
+    taz : pd.DataFrame
+        TAZ attributes with 'car_ownership_rate', 'x_km', 'y_km'.
+    travel_time : pd.DataFrame
+        Base car travel time matrix (minutes).
+    beta_time : float
+        Coefficient on in-vehicle travel time.
+    beta_cost : float
+        Coefficient on generalized cost.
+    beta_car_own : float
+        Additional utility for Car associated with car ownership rate.
+
+    Returns
+    -------
+    ModeChoiceResult
+        Aggregated OD, volumes by mode, and probabilities by mode.
+    """
+    zones = travel_time.index
+
+    # 1. Aggregate OD across purposes
+    total_od = None
+    for mat in od_matrices.values():
+        if total_od is None:
+            total_od = mat.copy()
+        else:
+            total_od += mat
+    total_od = total_od.loc[zones, zones]
+
+    # 2. Build time and cost matrices for each mode
+    tt_car = travel_time.loc[zones, zones].astype(float)
+    tt_metro = tt_car * 0.8  # metro faster
+    tt_bus = tt_car * 1.3    # bus slower
+
+    # Distance proxy (km)
+    dist_proxy = tt_car / 60.0 * 30.0  # 30 km/h
+
+    cost_car = 2.0 + 0.12 * dist_proxy
+    cost_metro = 15.0
+    cost_bus = 8.0 + 0.03 * dist_proxy
+
+    # 3. Car ownership matrix
+    car_own = taz["car_ownership_rate"].reindex(zones).to_numpy()
+    n = len(zones)
+    car_own_matrix = np.repeat(car_own[:, None], n, axis=1)
+
+    # 4. Utilities
+    modes = ["car", "metro", "bus"]
+    utilities = {}
+
+    # Car
+    U_car = (
+        beta_time * tt_car.to_numpy()
+        + beta_cost * cost_car.to_numpy()
+        + beta_car_own * car_own_matrix
+    )
+    utilities["car"] = U_car
+
+    # Metro
+    U_metro = beta_time * tt_metro.to_numpy() + beta_cost * cost_metro.to_numpy()
+    utilities["metro"] = U_metro
+
+    # Bus
+    U_bus = beta_time * tt_bus.to_numpy() + beta_cost * cost_bus.to_numpy()
+    utilities["bus"] = U_bus
+
+    # 5. Probabilities via softmax
+    exp_sum = np.zeros_like(U_car)
+    for U in utilities.values():
+        exp_sum += np.exp(U)
+
+    probabilities: Dict[str, pd.DataFrame] = {}
+    for mode, U in utilities.items():
+        P = np.exp(U) / np.maximum(exp_sum, 1e-12)
+        probabilities[mode] = pd.DataFrame(P, index=zones, columns=zones)
+
+    # 6. Mode-specific OD flows
+    volumes: Dict[str, pd.DataFrame] = {}
+    total_np = total_od.to_numpy()
+    for mode in modes:
+        volumes[mode] = pd.DataFrame(
+            total_np * probabilities[mode].to_numpy(),
+            index=zones,
+            columns=zones,
+        )
+
+    return ModeChoiceResult(
+        total_od=total_od,
+        volumes=volumes,
+        probabilities=probabilities,
+    )

modules/route_assignment.py

@@ -0,0 +1,216 @@
+# modules/route_assignment.py
+# TripAI – Synthetic Network + AON + Frank–Wolfe UE
+
+from __future__ import annotations
+from dataclasses import dataclass
+from typing import List
+
+import numpy as np
+import pandas as pd
+
+
+@dataclass
+class Network:
+    """
+    Simple synthetic network representation where each TAZ pair
+    (i, j) is connected by a single directed link.
+
+    Attributes
+    ----------
+    links : pd.DataFrame
+        Columns:
+        - link_id
+        - from_zone
+        - to_zone
+        - length_km
+        - t0_min        (free-flow travel time)
+        - capacity_vehph
+        - alpha         (BPR parameter)
+        - beta          (BPR parameter)
+    """
+    links: pd.DataFrame
+
+
+def generate_synthetic_network(taz: pd.DataFrame) -> Network:
+    """
+    Generate a fully connected directed network over TAZ centroids.
+    Each ordered pair (i, j), i != j, is represented as a distinct link.
+
+    Travel times are approximated from Euclidean distance and an
+    assumed free-flow speed.
+
+    Parameters
+    ----------
+    taz : pd.DataFrame
+        Must include 'x_km' and 'y_km' columns.
+
+    Returns
+    -------
+    Network
+    """
+    zones = taz.index.to_list()
+    coords = taz[["x_km", "y_km"]].to_numpy()
+    n = len(zones)
+
+    rows = []
+    link_id = 0
+    ff_speed_kmh = 30.0
+
+    for i_idx, i in enumerate(zones):
+        for j_idx, j in enumerate(zones):
+            if i == j:
+                continue
+            dx = coords[j_idx, 0] - coords[i_idx, 0]
+            dy = coords[j_idx, 1] - coords[i_idx, 1]
+            dist = np.sqrt(dx**2 + dy**2)  # km
+            t0 = (dist / max(ff_speed_kmh, 1e-3)) * 60.0 + 3.0  # minutes
+
+            rows.append(
+                {
+                    "link_id": link_id,
+                    "from_zone": i,
+                    "to_zone": j,
+                    "length_km": dist,
+                    "t0_min": t0,
+                    "capacity_vehph": np.random.uniform(1500, 2500),
+                    "alpha": 0.15,
+                    "beta": 4.0,
+                }
+            )
+            link_id += 1
+
+    links_df = pd.DataFrame(rows).set_index("link_id")
+    return Network(links=links_df)
+
+
+def _init_flow_column(links: pd.DataFrame, col: str = "flow_vehph") -> pd.DataFrame:
+    df = links.copy()
+    if col not in df.columns:
+        df[col] = 0.0
+    else:
+        df[col] = 0.0
+    return df
+
+
+def aon_assignment(od_car: pd.DataFrame, network: Network) -> pd.DataFrame:
+    """
+    All-or-nothing (AON) assignment assuming a single direct link
+    between each TAZ pair (i, j). All demand from i to j is loaded
+    on that link.
+
+    Parameters
+    ----------
+    od_car : pd.DataFrame
+        Car OD matrix (veh/h equivalent).
+    network : Network
+
+    Returns
+    -------
+    pd.DataFrame
+        Link flows with column 'flow_vehph'.
+    """
+    links = _init_flow_column(network.links, col="flow_vehph")
+    zones = od_car.index.to_list()
+
+    for i in zones:
+        for j in zones:
+            if i == j:
+                continue
+            q = float(od_car.loc[i, j])
+            if q <= 0:
+                continue
+            mask = (links["from_zone"] == i) & (links["to_zone"] == j)
+            links.loc[mask, "flow_vehph"] += q
+
+    return links
+
+
+def _bpr_travel_time(
+    flows: np.ndarray,
+    t0: np.ndarray,
+    capacity: np.ndarray,
+    alpha: np.ndarray,
+    beta: np.ndarray,
+) -> np.ndarray:
+    """Standard BPR volume-delay function."""
+    vc = np.divide(flows, capacity, out=np.zeros_like(flows), where=capacity > 0)
+    return t0 * (1.0 + alpha * np.power(vc, beta))
+
+
+def frank_wolfe_ue(
+    od_car: pd.DataFrame,
+    network: Network,
+    max_iter: int = 30,
+) -> pd.DataFrame:
+    """
+    Very simple Frank–Wolfe style User Equilibrium assignment over
+    the synthetic network where each OD pair has a single link.
+
+    Because there is only one 'route' per OD, the UE solution
+    coincides with the AON solution. This implementation still
+    outlines the iterative structure for pedagogical purposes.
+
+    Parameters
+    ----------
+    od_car : pd.DataFrame
+        Car OD matrix (veh/h).
+    network : Network
+    max_iter : int
+        Maximum iterations (for demonstration).
+
+    Returns
+    -------
+    pd.DataFrame
+        Link flows with column 'flow_vehph' and implied travel times.
+    """
+    links = network.links.copy()
+    n_links = len(links)
+
+    # Initialize flows
+    flows = np.zeros(n_links, dtype=float)
+
+    # Extract BPR parameters
+    t0 = links["t0_min"].to_numpy()
+    cap = links["capacity_vehph"].to_numpy()
+    alpha = links["alpha"].to_numpy()
+    beta = links["beta"].to_numpy()
+
+    # Pre-build a mapping (from_zone, to_zone) -> link indices
+    index = links.reset_index()
+    zone_pairs = {}
+    for idx, row in index.iterrows():
+        key = (row["from_zone"], row["to_zone"])
+        zone_pairs[key] = row["link_id"]
+
+    # Iterate Frank–Wolfe (though it converges immediately in this simple network)
+    for k in range(max_iter):
+        # Step 1: Compute travel times (not used for path choice here)
+        tt = _bpr_travel_time(flows, t0, cap, alpha, beta)
+
+        # Step 2: AON step (all or nothing given current times – here trivial)
+        aon_flows = np.zeros_like(flows)
+        zones = od_car.index.to_list()
+        for i in zones:
+            for j in zones:
+                if i == j:
+                    continue
+                q = float(od_car.loc[i, j])
+                if q <= 0:
+                    continue
+                lid = zone_pairs[(i, j)]
+                aon_flows[lid] += q
+
+        # Step 3: Line search step-size (generic diminishing rule)
+        step = 2.0 / (k + 2.0)
+        new_flows = flows + step * (aon_flows - flows)
+
+        # Convergence check
+        if np.allclose(new_flows, flows, atol=1e-3):
+            flows = new_flows
+            break
+
+        flows = new_flows
+
+    links["flow_vehph"] = flows
+    links["tt_min"] = _bpr_travel_time(flows, t0, cap, alpha, beta)
+    return links

modules/synthetic_city.py

@@ -0,0 +1,37 @@
+import streamlit as st
+import pandas as pd
+import os
+
+from modules.synthetic_city import generate_synthetic_city
+
+st.title("📊 Generate Synthetic City (20 TAZ)")
+
+# -------------------------------------
+# GENERATE SYNTHETIC CITY
+# -------------------------------------
+if st.button("Generate Synthetic Region"):
+    city = generate_synthetic_city()
+
+    # Save to session state
+    st.session_state["city"] = city
+    st.success("Synthetic city generated successfully!")
+
+# -------------------------------------
+# DISPLAY RESULTS
+# -------------------------------------
+if "city" in st.session_state:
+    city = st.session_state["city"]
+
+    st.subheader("TAZ Attributes")
+    st.dataframe(city.taz)
+
+    st.subheader("Summary Statistics")
+    st.write(city.taz.describe())
+
+    # Save files to /data/
+    os.makedirs("data", exist_ok=True)
+    city.taz.to_csv("data/taz_attributes.csv")
+    city.distance_matrix.to_csv("data/distance_matrix.csv")
+    city.travel_time_matrix.to_csv("data/travel_time_matrix.csv")
+
+    st.info("Files saved in /data/")

modules/trip_generation.py

@@ -0,0 +1,77 @@
+# modules/trip_generation.py
+# TripAI – Trip Generation Model
+
+from __future__ import annotations
+import pandas as pd
+import numpy as np
+
+PURPOSES = ["HBW", "HBE", "HBS"]
+
+
+def trip_generation(taz: pd.DataFrame) -> tuple[pd.DataFrame, pd.DataFrame]:
+    """
+    Compute trip productions and attractions for each TAZ for three purposes:
+    - HBW: Home–Based Work
+    - HBE: Home–Based Education
+    - HBS: Home–Based Shopping/Other
+
+    The functional forms are deliberately simple but grounded in standard
+    trip-rate logic and can be modified for calibration.
+
+    Parameters
+    ----------
+    taz : pd.DataFrame
+        TAZ-level attributes with at least the following columns:
+        ['households', 'workers', 'students', 'cars',
+         'service_jobs', 'industrial_jobs', 'retail_jobs',
+         'school_capacity', 'retail_floor_area'].
+
+    Returns
+    -------
+    productions : pd.DataFrame
+        Index = TAZ, columns = ['HBW', 'HBE', 'HBS'].
+    attractions : pd.DataFrame
+        Index = TAZ, columns = ['HBW', 'HBE', 'HBS'], balanced so that
+        sum(P) = sum(A) for each purpose.
+    """
+    df = taz.copy()
+
+    # ------------------------------------------------
+    # PRODUCTIONS (simple rate-based formulations)
+    # ------------------------------------------------
+    # HBW: mainly driven by workers and car availability
+    P_HBW = 0.8 * df["workers"] + 0.2 * df["cars"]
+
+    # HBE: driven by students
+    P_HBE = 1.2 * df["students"]
+
+    # HBS: driven by households (shopping, other)
+    P_HBS = 0.4 * df["households"]
+
+    productions = pd.DataFrame(
+        {"HBW": P_HBW, "HBE": P_HBE, "HBS": P_HBS},
+        index=df.index,
+    )
+
+    # ------------------------------------------------
+    # ATTRACTIONS (jobs, schools, retail)
+    # ------------------------------------------------
+    A_HBW = 0.7 * df["service_jobs"] + 0.3 * df["industrial_jobs"]
+    A_HBE = 1.5 * df["school_capacity"]
+    A_HBS = 1.3 * df["retail_floor_area"]
+
+    attractions = pd.DataFrame(
+        {"HBW": A_HBW, "HBE": A_HBE, "HBS": A_HBS},
+        index=df.index,
+    )
+
+    # ------------------------------------------------
+    # SIMPLE BALANCING (one-step scaling)
+    # ------------------------------------------------
+    for p in PURPOSES:
+        total_P = productions[p].sum()
+        total_A = attractions[p].sum()
+        if total_A > 0:
+            attractions[p] *= total_P / total_A
+
+    return productions, attractions

modules/utils.py

@@ -0,0 +1,64 @@
+# modules/utils.py
+# TripAI – Utility functions (IPF, etc.)
+
+from __future__ import annotations
+import numpy as np
+import pandas as pd
+
+
+def iterative_proportional_fitting(
+    T_init: np.ndarray,
+    row_targets: np.ndarray,
+    col_targets: np.ndarray,
+    max_iter: int = 50,
+    tol: float = 1e-6,
+) -> np.ndarray:
+    """
+    Perform Iterative Proportional Fitting (IPF) to adjust an initial
+    non-negative matrix T_init so that its row and column sums match
+    given marginals.
+
+    Parameters
+    ----------
+    T_init : np.ndarray
+        Initial non-negative matrix (NxN).
+    row_targets : np.ndarray
+        Target row sums (length N).
+    col_targets : np.ndarray
+        Target column sums (length N).
+    max_iter : int
+        Maximum number of IPF iterations.
+    tol : float
+        Convergence tolerance on row/column marginal differences.
+
+    Returns
+    -------
+    np.ndarray
+        Adjusted matrix with approximately matching row/column totals.
+    """
+    T = T_init.copy().astype(float)
+    n = T.shape[0]
+
+    for _ in range(max_iter):
+        # Row scaling
+        row_sums = T.sum(axis=1)
+        row_factors = np.ones(n)
+        mask = row_sums > 0
+        row_factors[mask] = row_targets[mask] / row_sums[mask]
+        T *= row_factors[:, None]
+
+        # Column scaling
+        col_sums = T.sum(axis=0)
+        col_factors = np.ones(n)
+        mask = col_sums > 0
+        col_factors[mask] = col_targets[mask] / col_sums[mask]
+        T *= col_factors[None, :]
+
+        # Check convergence
+        if (
+            np.allclose(T.sum(axis=1), row_targets, atol=tol)
+            and np.allclose(T.sum(axis=0), col_targets, atol=tol)
+        ):
+            break
+
+    return T