Align leave-one-out connected-set pruning with the official KSS MATLAB implementation

docs/make.jl

@@ -15,6 +15,7 @@ makedocs(;
         "Home" => "index.md"
         "Executable" => "Executable.md"
         "Package" => "Package.md"
+        "KSS pruning" => "Pruning.md"
     ],
 )
 

docs/src/Package.md

@@ -23,6 +23,10 @@ Depth = 3
 
 ## Main Function 
 
+See [KSS leave-one-worker-out pruning](Pruning.md) for the sample-construction
+algorithm and its distinction from the subsequent match- or observation-level
+bias correction.
+
 
 ```@docs
     leave_out_KSS(y,first_id,second_id;controls, do_lincom , Z_lincom , lincom_labels , settings)       
@@ -90,4 +94,3 @@ region[findall(region.==-1)].=0
 
 ```
 
-

docs/src/Pruning.md

@@ -0,0 +1,87 @@
+# KSS leave-one-worker-out pruning
+
+## What the pruning step guarantees
+
+Kline, Saggio, and Sølvsten (KSS) estimate variance components on a sample whose
+firm mobility network remains connected after removing any one worker's complete
+observation history. This is a property of the selected raw-observation sample.
+It is not a choice between leaving out a worker, a match, or an observation.
+
+The input must be sorted by worker and observation order, as required by
+`leave_out_KSS`. Pruning preserves original observation indices and returns
+dense worker and firm identifiers.
+
+The implementation follows the released MATLAB routines at commit
+[`8b957ff`](https://github.com/rsaggio87/LeaveOutTwoWay/tree/8b957ffeb10b8465a3584fceb0265cccc48379e1):
+
+- [`connected_set.m`](https://github.com/rsaggio87/LeaveOutTwoWay/blob/8b957ffeb10b8465a3584fceb0265cccc48379e1/codes/connected_set.m)
+  constructs the initial firm-transition graph.
+- [`pruning_unbal_v3.m`](https://github.com/rsaggio87/LeaveOutTwoWay/blob/8b957ffeb10b8465a3584fceb0265cccc48379e1/codes/pruning_unbal_v3.m)
+  removes articulation workers to a fixed point.
+- [`build_adj.m`](https://github.com/rsaggio87/LeaveOutTwoWay/blob/8b957ffeb10b8465a3584fceb0265cccc48379e1/codes/build_adj.m)
+  rebuilds the mover-induced firm-transition graph.
+- [`leave_out_KSS.m`](https://github.com/rsaggio87/LeaveOutTwoWay/blob/8b957ffeb10b8465a3584fceb0265cccc48379e1/codes/leave_out_KSS.m)
+  applies the final raw-observation singleton filter.
+
+The Julia implementation is an independent translation of that behavior.
+
+## Algorithm
+
+Given sorted raw observations `(worker, firm)`:
+
+1. Construct the firm-transition graph from consecutive observations within
+   workers and retain its unique largest component, ranked by number of firms.
+2. Identify movers and deduplicate their worker-firm pairs.
+3. Construct the mover-only bipartite graph and find articulation workers.
+4. Remove every articulation worker's complete raw-observation history.
+5. Rebuild the mover-induced firm-transition graph and retain its unique largest
+   component, again ranked by number of firms.
+6. Repeat steps 2--5 until no articulation worker remains.
+7. Drop workers represented by only one raw observation and relabel identifiers.
+
+The routine raises an `ArgumentError` if largest firm components are tied or if
+the mover graph disappears. Choosing a tied component by storage order would
+make the selected sample arbitrary.
+
+## Difference from the previous selector
+
+The previous Julia selector built the full worker-firm bipartite graph and ranked
+components by their total number of worker and firm vertices. The released
+MATLAB code ranks components of the firm-transition graph by number of firms.
+Those objectives can select different surviving samples.
+
+The graph used to detect articulation workers is not the source of this
+difference. Stayers are degree-one leaves. A leaf is never a cut vertex, and
+attaching leaves to firm vertices cannot change whether a worker vertex is a cut
+vertex. Consequently, the full and mover-only bipartite graphs identify the same
+articulation workers on a common sample. They can diverge when the components
+remaining after deletion are ranked.
+
+For example, consider a two-worker, three-firm `K(2, 3)` mover core joined by an
+articulation worker to one firm with six two-observation stayers. After removing
+the bridge worker, firm ranking retains the three-firm mover core (six
+observations), while total-vertex ranking retains the one-firm stayer component
+(twelve observations).
+
+The previous routine also retained a stale worker count across iterations. This
+added isolated phantom vertices and wasted memory, but it does not appear to
+alter articulation workers or component selection. Graph dimensions are now
+recomputed from the current sample in every round.
+
+Both selectors happen to agree on the public `test.csv` fixture:
+
+- 71,614 input observations;
+- 63,462 observations after the initial firm component;
+- 56,044 final observations;
+- 28,022 workers, 1,684 firms, 34,436 matches, and 6,414 movers.
+
+Agreement on this fixture does not establish that the algorithms are equivalent;
+the synthetic component-ranking example demonstrates the difference.
+
+## Pruning versus `leave_out_level`
+
+`get_leave_one_out_set` always constructs the leave-one-worker-out-connected raw
+observation sample described above. The `VCHDFESettings.leave_out_level` value
+does not change pruning. Its `"match"` and `"obs"` settings apply only
+later in `leave_out_estimation`, when the package computes the KSS bias
+correction on the already selected sample.

docs/src/index.md

@@ -10,6 +10,6 @@ Please report any bug or problem [here](https://github.com/HighDimensionalEconLa
 # Contents
 
 ```@contents
-Pages = ["Executable.md", "Package.md"]
+Pages = ["Executable.md", "Package.md", "Pruning.md"]
 Depth = 3
-```
+```

src/VarianceComponentsHDFE.jl

@@ -11,6 +11,7 @@ using ArgParse
 using FixedEffects
 
 #include("init.jl")
+include("kss_pruning.jl")
 include("leave_out_correction.jl")
 include("parameters_settings.jl")
 include("laplacians/Laplacians.jl")

src/kss_pruning.jl

@@ -0,0 +1,191 @@
+function _check_pruning_inputs(y, first_id, second_id, obs_id = nothing)
+    n_observations = length(y)
+    length(first_id) == n_observations ||
+        throw(ArgumentError("y and first_id must have the same length"))
+    length(second_id) == n_observations ||
+        throw(ArgumentError("y and second_id must have the same length"))
+    obs_id === nothing || length(obs_id) == n_observations ||
+        throw(ArgumentError("y and obs_id must have the same length"))
+    n_observations > 0 ||
+        throw(ArgumentError("KSS pruning requires at least one observation"))
+    return nothing
+end
+
+function _dense_ids(ids)
+    levels = unique(sort(ids))
+    return Int.(indexin(ids, levels)), levels
+end
+
+function _select_and_relabel(y, first_id, second_id, obs_id, selection)
+    selected_y = y[selection]
+    selected_first_id = first_id[selection]
+    selected_second_id = second_id[selection]
+    selected_obs_id = obs_id[selection]
+
+    dense_first_id, _ = _dense_ids(selected_first_id)
+    dense_second_id, _ = _dense_ids(selected_second_id)
+    return (
+        obs_id = selected_obs_id,
+        y = selected_y,
+        first_id = dense_first_id,
+        second_id = dense_second_id,
+    )
+end
+
+function _firm_transition_graph(first_id, second_id, firms = unique(sort(second_id)))
+    firm_id = Int.(indexin(second_id, firms))
+    graph = Graph(length(firms))
+    for observation = 2:length(first_id)
+        same_worker = first_id[observation] == first_id[observation - 1]
+        changed_firm = firm_id[observation] != firm_id[observation - 1]
+        same_worker && changed_firm &&
+            add_edge!(graph, firm_id[observation - 1], firm_id[observation])
+    end
+    return graph, firms
+end
+
+function _unique_largest_firm_component(graph, stage)
+    components = connected_components(graph)
+    isempty(components) &&
+        throw(ArgumentError("KSS pruning found no firm component $stage"))
+
+    component_sizes = length.(components)
+    largest_size = maximum(component_sizes)
+    largest_components = findall(==(largest_size), component_sizes)
+    if length(largest_components) != 1
+        throw(
+            ArgumentError(
+                "KSS pruning found $(length(largest_components)) tied largest firm " *
+                "components $stage ($largest_size firms each)",
+            ),
+        )
+    end
+    return components[only(largest_components)]
+end
+
+function _mover_workers(first_id, second_id)
+    first_firm = zeros(Int, maximum(first_id))
+    movers = falses(length(first_firm))
+    for observation in eachindex(first_id)
+        worker = first_id[observation]
+        firm = second_id[observation]
+        if first_firm[worker] == 0
+            first_firm[worker] = firm
+        elseif first_firm[worker] != firm
+            movers[worker] = true
+        end
+    end
+    return findall(movers)
+end
+
+function _articulation_workers(first_id, second_id)
+    mover_workers = _mover_workers(first_id, second_id)
+    isempty(mover_workers) && throw(
+        ArgumentError(
+            "KSS pruning cannot continue because the mover graph has disappeared",
+        ),
+    )
+
+    mover_index = zeros(Int, maximum(first_id))
+    for (index, worker) in enumerate(mover_workers)
+        mover_index[worker] = index
+    end
+    mover_rows = findall(worker -> mover_index[worker] > 0, first_id)
+    pairs = Tuple{Int,Int}[]
+    for row in mover_rows
+        push!(pairs, (mover_index[first_id[row]], second_id[row]))
+    end
+    unique!(pairs)
+
+    n_movers = length(mover_workers)
+    graph = Graph(n_movers + maximum(second_id))
+    for (worker, firm) in pairs
+        add_edge!(graph, worker, n_movers + firm)
+    end
+
+    articulation_workers = filter(vertex -> vertex <= n_movers, articulation(graph))
+    return mover_workers[articulation_workers]
+end
+
+function _largest_mover_firm_sample(y, first_id, second_id, obs_id)
+    mover_workers = _mover_workers(first_id, second_id)
+    isempty(mover_workers) && throw(
+        ArgumentError(
+            "KSS pruning cannot continue because the mover graph has disappeared",
+        ),
+    )
+
+    is_mover = falses(maximum(first_id))
+    is_mover[mover_workers] .= true
+    mover_rows = findall(worker -> is_mover[worker], first_id)
+    firms = collect(1:maximum(second_id))
+    graph, firms = _firm_transition_graph(
+        first_id[mover_rows],
+        second_id[mover_rows],
+        firms,
+    )
+    component = _unique_largest_firm_component(graph, "after articulation deletion")
+    connected_firms = Set(firms[component])
+    selection = findall(firm -> firm in connected_firms, second_id)
+    return _select_and_relabel(y, first_id, second_id, obs_id, selection)
+end
+
+function _official_find_connected_set(y, first_id, second_id, settings)
+    _check_pruning_inputs(y, first_id, second_id)
+    obs_id = collect(1:length(y))
+    graph, firms = _firm_transition_graph(first_id, second_id)
+    component = _unique_largest_firm_component(graph, "in the initial sample")
+    connected_firms = Set(firms[component])
+    selection = findall(firm -> firm in connected_firms, second_id)
+    return _select_and_relabel(y, first_id, second_id, obs_id, selection)
+end
+
+function _official_prunning_connected_set(
+    y,
+    first_idvar,
+    second_idvar,
+    obs_id,
+    settings,
+)
+    _check_pruning_inputs(y, first_idvar, second_idvar, obs_id)
+    first_id, _ = _dense_ids(first_idvar)
+    second_id, _ = _dense_ids(second_idvar)
+
+    while true
+        bad_workers = _articulation_workers(first_id, second_id)
+        n_bad_workers = length(bad_workers)
+        settings.print_level > 1 && println(
+            "Number of $(settings.first_id_display_small) that are articulation points: ",
+            n_bad_workers,
+        )
+        n_bad_workers == 0 && break
+
+        is_bad_worker = falses(maximum(first_id))
+        is_bad_worker[bad_workers] .= true
+        selection = findall(worker -> !is_bad_worker[worker], first_id)
+        isempty(selection) && throw(
+            ArgumentError(
+                "KSS pruning removed all observations before reaching a fixed point",
+            ),
+        )
+
+        sample = _select_and_relabel(y, first_id, second_id, obs_id, selection)
+        sample = _largest_mover_firm_sample(
+            sample.y,
+            sample.first_id,
+            sample.second_id,
+            sample.obs_id,
+        )
+        obs_id = sample.obs_id
+        y = sample.y
+        first_id = sample.first_id
+        second_id = sample.second_id
+    end
+
+    return (
+        obs_id = obs_id,
+        y = y,
+        first_id = first_id,
+        second_id = second_id,
+    )
+end