import pandas as pd

import larch as lx
from larch.util.namespaces import Namespace
spec_content = """expression,DA,SR,Walk,Bike,Transit
AUTO_TIME,InVehTime,InVehTime,,,
AUTO_COST,Cost,,,,
AUTO_COST*0.5,,Cost,,,
1,,ASC_SR,ASC_Walk,ASC_Bike,ASC_Transit
log(INCOME),,LogIncome_SR,LogIncome_Walk,LogIncome_Bike,LogIncome_Transit
WALK_TIME,,,NonMotorTime,,
BIKE_TIME,,,,NonMotorTime,
TRANSIT_IVTT,,,,,InVehTime
TRANSIT_OVTT,,,,,OutVehTime
TRANSIT_FARE,,,,,Cost
"""

from io import StringIO

spec = pd.read_csv(StringIO(spec_content)).set_index("expression")
spec
DA SR Walk Bike Transit
expression
AUTO_TIME InVehTime InVehTime NaN NaN NaN
AUTO_COST Cost NaN NaN NaN NaN
AUTO_COST*0.5 NaN Cost NaN NaN NaN
1 NaN ASC_SR ASC_Walk ASC_Bike ASC_Transit
log(INCOME) NaN LogIncome_SR LogIncome_Walk LogIncome_Bike LogIncome_Transit
WALK_TIME NaN NaN NonMotorTime NaN NaN
BIKE_TIME NaN NaN NaN NonMotorTime NaN
TRANSIT_IVTT NaN NaN NaN NaN InVehTime
TRANSIT_OVTT NaN NaN NaN NaN OutVehTime
TRANSIT_FARE NaN NaN NaN NaN Cost
hh, pp, tour, skims = lx.example(200, ["hh", "pp", "tour", "skims"])
Mode = Namespace(
    DA=1,
    SR=2,
    Walk=3,
    Bike=4,
    Transit=5,
)
# tour_dataset = lx.Dataset.construct.from_idco(tour.set_index('TOURID'), alts=Mode)

tour_dataset = tour.set_index("TOURID")
od_skims = lx.Dataset.construct.from_omx(skims)
dt = lx.DataTree(
    tour=tour_dataset,
    hh=hh.set_index("HHID"),
    person=pp.set_index("PERSONID"),
    od=od_skims,
    do=od_skims,
    relationships=(
        "tours.HHID @ hh.HHID",
        "tours.PERSONID @ person.PERSONID",
        "hh.HOMETAZ @ od.otaz",
        "tours.DTAZ @ od.dtaz",
        "hh.HOMETAZ @ do.dtaz",
        "tours.DTAZ @ do.otaz",
    ),
    root_node_name="tour",
)
spec.columns
Index(['DA', 'SR', 'Walk', 'Bike', 'Transit'], dtype='object')
dt.set_altnames(spec.columns)
dt.alts_name_to_id()
{'DA': np.int64(1),
 'SR': np.int64(2),
 'Walk': np.int64(3),
 'Bike': np.int64(4),
 'Transit': np.int64(5)}
import larch.model.from_spec
m = lx.Model.from_spec(spec, dt)
m.choice_co_code = "TOURMODE"
m.should_preload_data()
m.maximize_loglike()
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓
┃ Larch Model Dashboard                                                                 ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩
│ optimization complete                                                                 │
│ Log Likelihood Current =     -11021.349609 Best =     -11021.349609                   │
│ ┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━┓ │
│ ┃ Parameter                       Value           Gradient        Best           ┃ │
│ ┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━┩ │
│ │ ASC_Bike                       │  0.6916911     │ -0.009673357   │  0.6916911     │ │
│ │ ASC_SR                         │  3.1878242     │  0.058359265   │  3.1878242     │ │
│ │ ASC_Transit                    │  2.8104834     │  0.017008305   │  2.8104834     │ │
│ │ ASC_Walk                       │  6.7744329     │  0.030028343   │  6.7744329     │ │
│ │ Cost                           │  0.5784966     │ -0.033325195   │  0.5784966     │ │
│ │ InVehTime                      │ -0.20312673    │ -0.30273438    │ -0.20312673    │ │
│ │ LogIncome_Bike                 │ -0.30356777    │ -0.10908508    │ -0.30356777    │ │
│ │ LogIncome_SR                   │ -0.43742301    │  0.61865234    │ -0.43742301    │ │
│ │ LogIncome_Transit              │ -0.37474969    │  0.16412354    │ -0.37474969    │ │
│ │ LogIncome_Walk                 │ -0.36612581    │  0.31430054    │ -0.36612581    │ │
│ │ NonMotorTime                   │ -0.24666207    │  0.73594666    │ -0.24666207    │ │
│ │ OutVehTime                     │ -0.27965134    │  0.23824883    │ -0.27965134    │ │
│ └────────────────────────────────┴────────────────┴────────────────┴────────────────┘ │
└───────────────────────────────────────────────────────────────────────────────────────┘
      message: Optimization terminated successfully
      success: True
       status: 0
            x: [ 6.917e-01  3.188e+00  2.810e+00  6.774e+00  5.785e-01
                -2.031e-01 -3.036e-01 -4.374e-01 -3.747e-01 -3.661e-01
                -2.467e-01 -2.797e-01]
          nit: 75
          jac: [-9.673e-03  5.836e-02  1.701e-02  3.003e-02 -3.333e-02
                -3.027e-01 -1.091e-01  6.187e-01  1.641e-01  3.143e-01
                 7.359e-01  2.382e-01]
         nfev: 177
         njev: 75
      n_cases: 20739
 total_weight: 20739.0
      logloss: 0.5314311012765803
      loglike: -11021.349609375
dt.n_cases
20739
m.logsums()
array([ 0.09595297, -1.95442336, -1.94657884, ..., -0.44695921,
       -0.56635064, -0.13907752], shape=(20739,))
m.pvals = {"Cost": -0.1, "InVehTime": -0.2, "ASC_Bike": 2.0}
m.pf
value best initvalue minimum maximum nullvalue holdfast
param_name
ASC_Bike 2.000000 0.691691 0.0 -inf inf 0.0 0
ASC_SR 3.187824 3.187824 0.0 -inf inf 0.0 0
ASC_Transit 2.810483 2.810483 0.0 -inf inf 0.0 0
ASC_Walk 6.774433 6.774433 0.0 -inf inf 0.0 0
Cost -0.100000 0.578497 0.0 -inf inf 0.0 0
InVehTime -0.200000 -0.203127 0.0 -inf inf 0.0 0
LogIncome_Bike -0.303568 -0.303568 0.0 -inf inf 0.0 0
LogIncome_SR -0.437423 -0.437423 0.0 -inf inf 0.0 0
LogIncome_Transit -0.374750 -0.374750 0.0 -inf inf 0.0 0
LogIncome_Walk -0.366126 -0.366126 0.0 -inf inf 0.0 0
NonMotorTime -0.246662 -0.246662 0.0 -inf inf 0.0 0
OutVehTime -0.279651 -0.279651 0.0 -inf inf 0.0 0