Example usage for java.util TreeMap compute

Introduction

In this page you can find the example usage for java.util TreeMap compute.

Prototype

default V compute(K key, BiFunction<? super K, ? super V, ? extends V> remappingFunction) 

Source Link

Document

Attempts to compute a mapping for the specified key and its current mapped value (or null if there is no current mapping).

Usage

From source file:nl.rivm.cib.episim.model.disease.infection.MSEIRSTest.java

public static Observable<Entry<Double, long[]>> stochasticSellke(final SIRConfig config, final double maxDt) {
    return Observable.create(sub -> {
        final double beta = config.reproduction() / config.recovery();
        final long[] y = config.population();
        final double[] T = config.t();
        final double dt = Double.isFinite(maxDt) && maxDt > 0 ? maxDt : T[1];

        final Long seed = config.seed();
        final RandomGenerator rng = new MersenneTwister(seed == null ? System.currentTimeMillis() : seed);

        final ExponentialDistribution resistanceDist = new ExponentialDistribution(rng, 1),
                recoverDist = new ExponentialDistribution(rng, config.recovery());

        // pending infections (mapping resistance -> amount)
        final TreeMap<Double, Integer> tInfect = IntStream.range(0, (int) y[0])
                .mapToObj(i -> resistanceDist.sample())
                .collect(Collectors.toMap(r -> r, r -> 1, Integer::sum, TreeMap::new));
        // pending recoveries (mapping time -> amount)
        final TreeMap<Double, Integer> tRecover = new TreeMap<>();

        double cumPres = 0;
        // Re-initialize infectives as susceptibles with zero resistance
        tInfect.put(cumPres, (int) y[1]);
        y[0] += y[1]; // I -> S
        y[1] -= y[1]; // I -> 0
        for (double t = T[0]; t < T[1];) {
            publishCopy(sub, t, y);/* w  w  w .j  av  a 2 s  .c  o  m*/
            final long localPopSize = y[0] + y[1] + y[2];
            final Double ri = tInfect.isEmpty() ? null : tInfect.firstKey(), ti = ri == null ? null :
            // now + remaining resistance per relative pressure
            t + (ri - cumPres) / (beta * Math.max(y[1], 1) / localPopSize),
                    tr = tRecover.isEmpty() ? null : tRecover.firstKey();

            // time of next infection is earliest
            if (ti != null && (tr == null || ti < tr)) {
                final int ni = tInfect.remove(ri);
                cumPres = ri;

                // publish intermediate values
                for (double t1 = Math.min(ti, t + dt), tMax = Math.min(T[1], ti); t1 < tMax; t1 += dt)
                    publishCopy(sub, t1, y);

                // infect
                t = ti;
                y[0] -= ni; // from S
                y[1] += ni; // to I

                // schedule S_t recoveries at t+Exp(1/gamma)
                for (int i = 0; i < ni; i++)
                    tRecover.compute(t + recoverDist.sample(), (k, v) -> v == null ? 1 : v + 1);
            }
            // time of next recovery is earliest
            else if (tr != null) {
                final int nr = tRecover.remove(tr);
                if (ri != null)
                    // advance cumulative pressure by dt * relative pressure
                    cumPres += (tr - t) * beta * y[1] / localPopSize;

                // publish intermediate values
                for (double t1 = Math.min(tr, t + dt), tMax = Math.min(T[1], tr); t1 < tMax; t1 += dt)
                    publishCopy(sub, t1, y);

                // recover
                t = tr;
                y[1] -= nr; // from I
                y[2] += nr; // to R
            }
            // no events remaining
            else {
                // publish intermediate values
                for (double t1 = t + dt; t1 < T[1]; t1 += dt)
                    publishCopy(sub, t1, y);

                // time ends
                break;
            }
        }
        sub.onComplete();
    });
}

---