zog

// this code is autogenerated, do not check it into to your code repository

// ref lib/zog/community.ex:21
const std = @import("std");
const beam = @import("beam");
const zog = @import("zog");

const ArrayGraph = zog.models.ArrayGraph;

fn buildGraph(node_count: usize, from: []u32, to: []u32, weight: []f64) !ArrayGraph(void, f64) {
    const allocator = beam.allocator;
    var g = ArrayGraph(void, f64).init(allocator);
    errdefer g.deinit();

    try g.nodes.ensureTotalCapacity(allocator, node_count);
    try g.edges.ensureTotalCapacity(allocator, from.len);

    for (0..node_count) |_| {
        _ = try g.addNode({});
    }

    for (from, to, weight) |f, t, w| {
        _ = try g.addEdge(f, t, w);
    }

    return g;
}

fn extractAssignments(result: anytype, node_count: usize) ![]usize {
    const allocator = beam.allocator;
    var assignments = try allocator.alloc(usize, node_count);
    errdefer allocator.free(assignments);

    for (0..node_count) |i| {
        assignments[i] = result.assignments.get(@intCast(i)) orelse 0;
    }

    return assignments;
}

pub fn louvain(
    node_count: usize,
    from: []u32,
    to: []u32,
    weight: []f64,
    min_modularity_gain: f64,
    max_iterations: usize,
    seed: u64,
) ![]usize {
    var g = try buildGraph(node_count, from, to, weight);
    defer g.deinit();

    var result = try zog.community.louvain.detectWeightedWithOptions(
        beam.allocator,
        g,
        .{
            .min_modularity_gain = min_modularity_gain,
            .max_iterations = max_iterations,
            .seed = seed,
        },
        zog.utils.identityF64,
    );
    defer result.deinit();

    return extractAssignments(result, node_count);
}

pub fn leiden(
    node_count: usize,
    from: []u32,
    to: []u32,
    weight: []f64,
    min_modularity_gain: f64,
    max_iterations: usize,
    seed: u64,
    theta: f64,
) ![]usize {
    var g = try buildGraph(node_count, from, to, weight);
    defer g.deinit();

    var result = try zog.community.leiden.detectWeightedWithOptions(
        beam.allocator,
        g,
        .{
            .min_modularity_gain = min_modularity_gain,
            .max_iterations = max_iterations,
            .seed = seed,
            .theta = theta,
        },
        zog.utils.identityF64,
    );
    defer result.deinit();

    return extractAssignments(result, node_count);
}

pub fn leiden_hierarchical(
    node_count: usize,
    from: []u32,
    to: []u32,
    weight: []f64,
    min_modularity_gain: f64,
    max_iterations: usize,
    seed: u64,
    theta: f64,
) ![][]usize {
    var g = try buildGraph(node_count, from, to, weight);
    defer g.deinit();

    var result = try zog.community.leiden.detectHierarchicalWeightedWithOptions(
        beam.allocator,
        g,
        .{
            .min_modularity_gain = min_modularity_gain,
            .max_iterations = max_iterations,
            .seed = seed,
            .theta = theta,
        },
        zog.utils.identityF64,
    );
    defer result.deinit();

    const allocator = beam.allocator;
    const outer = try allocator.alloc([]usize, result.levels.len);
    errdefer allocator.free(outer);

    for (result.levels, 0..) |level, i| {
        const level_copy = try allocator.alloc(usize, node_count);
        errdefer allocator.free(level_copy);
        @memcpy(level_copy, level);
        outer[i] = level_copy;
    }

    return outer;
}

pub fn label_propagation(
    node_count: usize,
    from: []u32,
    to: []u32,
    weight: []f64,
    max_iterations: usize,
    seed: u64,
) ![]usize {
    var g = try buildGraph(node_count, from, to, weight);
    defer g.deinit();

    var result = try zog.community.label_propagation.labelPropagation(
        beam.allocator,
        g,
        .{
            .max_iterations = max_iterations,
            .seed = seed,
        },
    );
    defer result.deinit();

    return extractAssignments(result, node_count);
}

pub fn modularity_f64(
    node_count: usize,
    from: []u32,
    to: []u32,
    weight: []f64,
    assignments: []usize,
) !f64 {
    var g = try buildGraph(node_count, from, to, weight);
    defer g.deinit();

    var map = std.AutoHashMap(u32, usize).init(beam.allocator);
    defer map.deinit();

    for (assignments, 0..) |comm, i| {
        try map.put(@intCast(i), comm);
    }

    return try zog.community.metrics.modularity(beam.allocator, g, map, zog.utils.identityF64);
}