diff --git a/madspace/CMakeLists.txt b/madspace/CMakeLists.txt index d8231d185f..b16193713d 100644 --- a/madspace/CMakeLists.txt +++ b/madspace/CMakeLists.txt @@ -137,6 +137,7 @@ add_library( src/phasespace/luminosity.cpp src/phasespace/topology.cpp src/phasespace/t_propagator_mapping.cpp + src/phasespace/color_ordered_mapping.cpp src/phasespace/phasespace.cpp src/phasespace/multichannel.cpp src/phasespace/madnis.cpp @@ -189,6 +190,7 @@ add_library( include/madspace/phasespace/luminosity.hpp include/madspace/phasespace/topology.hpp include/madspace/phasespace/t_propagator_mapping.hpp + include/madspace/phasespace/color_ordered_mapping.hpp include/madspace/phasespace/phasespace.hpp include/madspace/phasespace/multichannel.hpp include/madspace/phasespace/madnis.hpp diff --git a/madspace/generate_pyi.sh b/madspace/generate_pyi.sh index 82bc3082c6..a1b009f8d3 100755 --- a/madspace/generate_pyi.sh +++ b/madspace/generate_pyi.sh @@ -2,4 +2,5 @@ pybind11-stubgen madspace._madspace_py \ -o . \ --enum-class-locations=TChannelMode:PhaseSpaceMapping.TChannelMode \ --enum-class-locations=Activation:MLP.Activation \ - --enum-class-locations=CutMode:Cuts.CutMode + --enum-class-locations=CutMode:Cuts.CutMode \ + --enum-class-locations=LRSchedule:AdamOptimizer.LRSchedule diff --git a/madspace/include/madspace/compgraphs/function_builder_mixin.inc b/madspace/include/madspace/compgraphs/function_builder_mixin.inc index ae220de13c..b1050b404b 100644 --- a/madspace/include/madspace/compgraphs/function_builder_mixin.inc +++ b/madspace/include/madspace/compgraphs/function_builder_mixin.inc @@ -87,12 +87,12 @@ Value reduce_sum_vector(Value in1) { return instruction("reduce_sum_vector", {in1})[0]; } -Value batch_reduce_mean(Value in) { - return instruction("batch_reduce_mean", {in})[0]; +Value batch_reduce_mean(Value in1) { + return instruction("batch_reduce_mean", {in1})[0]; } -Value batch_reduce_mean_keepdim(Value in) { - return instruction("batch_reduce_mean_keepdim", {in})[0]; +Value batch_reduce_mean_keepdim(Value in1) { + return instruction("batch_reduce_mean_keepdim", {in1})[0]; } Value reduce_product(Value in1) { @@ -260,6 +260,16 @@ std::array two_to_three_particle_scattering_inverse(Value p1, Value p2 return {output_vector[0], output_vector[1], output_vector[2], output_vector[3]}; } +std::array double_t_scattering(Value r_phi, Value pa, Value pb, Value t1_abs, Value t2_abs, Value m1) { + auto output_vector = instruction("double_t_scattering", {r_phi, pa, pb, t1_abs, t2_abs, m1}); + return {output_vector[0], output_vector[1], output_vector[2]}; +} + +std::array double_t_scattering_inverse(Value p1, Value p2, Value pa, Value pb) { + auto output_vector = instruction("double_t_scattering_inverse", {p1, p2, pa, pb}); + return {output_vector[0], output_vector[1]}; +} + std::array three_body_decay_com(Value r_e1, Value r_e2, Value r_phi, Value r_cos_theta, Value r_beta, Value m0, Value m1, Value m2, Value m3) { auto output_vector = instruction("three_body_decay_com", {r_e1, r_e2, r_phi, r_cos_theta, r_beta, m0, m1, m2, m3}); return {output_vector[0], output_vector[1], output_vector[2], output_vector[3]}; @@ -290,6 +300,26 @@ std::array t_inv_value_and_min_max(Value pa, Value pb, Value p1, Value return {output_vector[0], output_vector[1], output_vector[2]}; } +std::array t1_inv_min_max_doublet(Value pa, Value pb, Value m1, Value mir_min) { + auto output_vector = instruction("t1_inv_min_max_doublet", {pa, pb, m1, mir_min}); + return {output_vector[0], output_vector[1]}; +} + +std::array t1_inv_value_and_min_max_doublet(Value pa, Value pb, Value p1, Value m1, Value mir_min) { + auto output_vector = instruction("t1_inv_value_and_min_max_doublet", {pa, pb, p1, m1, mir_min}); + return {output_vector[0], output_vector[1], output_vector[2]}; +} + +std::array t2_inv_min_max_doublet(Value pa, Value pb, Value m1, Value mir_min, Value t1_abs) { + auto output_vector = instruction("t2_inv_min_max_doublet", {pa, pb, m1, mir_min, t1_abs}); + return {output_vector[0], output_vector[1]}; +} + +std::array t2_inv_value_and_min_max_doublet(Value pa, Value pb, Value p1, Value m1, Value mir_min, Value t1_abs) { + auto output_vector = instruction("t2_inv_value_and_min_max_doublet", {pa, pb, p1, m1, mir_min, t1_abs}); + return {output_vector[0], output_vector[1], output_vector[2]}; +} + std::array s23_min_max(Value pa, Value pb, Value p3, Value t1_abs, Value m1, Value m2) { auto output_vector = instruction("s23_min_max", {pa, pb, p3, t1_abs, m1, m2}); return {output_vector[0], output_vector[1]}; diff --git a/madspace/include/madspace/compgraphs/opcode_mixin.inc b/madspace/include/madspace/compgraphs/opcode_mixin.inc index 58a2cd2294..39c4dadb3d 100644 --- a/madspace/include/madspace/compgraphs/opcode_mixin.inc +++ b/madspace/include/madspace/compgraphs/opcode_mixin.inc @@ -59,86 +59,92 @@ two_to_two_particle_scattering = 57, two_to_two_particle_scattering_inverse = 58, two_to_three_particle_scattering = 59, two_to_three_particle_scattering_inverse = 60, -three_body_decay_com = 61, -three_body_decay_com_inverse = 62, -three_body_decay = 63, -three_body_decay_inverse = 64, -t_inv_min_max = 65, -t_inv_value_and_min_max = 66, -s23_min_max = 67, -s23_value_and_min_max = 68, -invariants_from_momenta = 69, -sde2_channel_weights = 70, -subchannel_weights = 71, -apply_subchannel_weights = 72, -pt_eta_phi_x = 73, -mirror_momenta = 74, -momenta_to_x1x2 = 75, -uniform_invariant = 76, -uniform_invariant_inverse = 77, -breit_wigner_invariant = 78, -breit_wigner_invariant_inverse = 79, -stable_invariant = 80, -stable_invariant_inverse = 81, -stable_invariant_nu = 82, -stable_invariant_nu_inverse = 83, -fast_rambo_massless = 84, -fast_rambo_massless_inverse = 85, -fast_rambo_massless_com = 86, -fast_rambo_massive = 87, -fast_rambo_massive_inverse = 88, -fast_rambo_massive_com = 89, -cut_unphysical = 90, -cut_one = 91, -cut_all = 92, -cut_any = 93, -scale_transverse_energy = 94, -scale_transverse_mass = 95, -scale_half_transverse_mass = 96, -scale_partonic_energy = 97, -chili_forward = 98, -chili_inverse = 99, -matrix_element = 100, -collect_channel_weights = 101, -interpolate_pdf = 102, -interpolate_alpha_s = 103, -matmul = 104, -relu = 105, -leaky_relu = 106, -elu = 107, -gelu = 108, -sigmoid = 109, -softplus = 110, -rqs_reshape = 111, -rqs_find_bin = 112, -rqs_forward = 113, -rqs_inverse = 114, -softmax = 115, -softmax_prior = 116, -sample_discrete = 117, -sample_discrete_inverse = 118, -sample_discrete_probs = 119, -sample_discrete_probs_inverse = 120, -discrete_histogram = 121, -permute_momenta = 122, -gather = 123, -gather_int = 124, -select_int = 125, -select = 126, -select_vector = 127, -argsort = 128, -one_hot = 129, -madnis_abs_weight = 130, -madnis_softclip = 131, -madnis_variance = 132, -madnis_single_channel_variance = 133, -madnis_multi_channel_variance = 134, -nonzero = 135, -batch_gather = 136, -batch_scatter = 137, -random = 138, -unweight = 139, -vegas_forward = 140, -vegas_inverse = 141, -vegas_histogram = 142, -histogram = 143 +double_t_scattering = 61, +double_t_scattering_inverse = 62, +three_body_decay_com = 63, +three_body_decay_com_inverse = 64, +three_body_decay = 65, +three_body_decay_inverse = 66, +t_inv_min_max = 67, +t_inv_value_and_min_max = 68, +t1_inv_min_max_doublet = 69, +t1_inv_value_and_min_max_doublet = 70, +t2_inv_min_max_doublet = 71, +t2_inv_value_and_min_max_doublet = 72, +s23_min_max = 73, +s23_value_and_min_max = 74, +invariants_from_momenta = 75, +sde2_channel_weights = 76, +subchannel_weights = 77, +apply_subchannel_weights = 78, +pt_eta_phi_x = 79, +mirror_momenta = 80, +momenta_to_x1x2 = 81, +uniform_invariant = 82, +uniform_invariant_inverse = 83, +breit_wigner_invariant = 84, +breit_wigner_invariant_inverse = 85, +stable_invariant = 86, +stable_invariant_inverse = 87, +stable_invariant_nu = 88, +stable_invariant_nu_inverse = 89, +fast_rambo_massless = 90, +fast_rambo_massless_inverse = 91, +fast_rambo_massless_com = 92, +fast_rambo_massive = 93, +fast_rambo_massive_inverse = 94, +fast_rambo_massive_com = 95, +cut_unphysical = 96, +cut_one = 97, +cut_all = 98, +cut_any = 99, +scale_transverse_energy = 100, +scale_transverse_mass = 101, +scale_half_transverse_mass = 102, +scale_partonic_energy = 103, +chili_forward = 104, +chili_inverse = 105, +matrix_element = 106, +collect_channel_weights = 107, +interpolate_pdf = 108, +interpolate_alpha_s = 109, +matmul = 110, +relu = 111, +leaky_relu = 112, +elu = 113, +gelu = 114, +sigmoid = 115, +softplus = 116, +rqs_reshape = 117, +rqs_find_bin = 118, +rqs_forward = 119, +rqs_inverse = 120, +softmax = 121, +softmax_prior = 122, +sample_discrete = 123, +sample_discrete_inverse = 124, +sample_discrete_probs = 125, +sample_discrete_probs_inverse = 126, +discrete_histogram = 127, +permute_momenta = 128, +gather = 129, +gather_int = 130, +select_int = 131, +select = 132, +select_vector = 133, +argsort = 134, +one_hot = 135, +madnis_abs_weight = 136, +madnis_softclip = 137, +madnis_variance = 138, +madnis_single_channel_variance = 139, +madnis_multi_channel_variance = 140, +nonzero = 141, +batch_gather = 142, +batch_scatter = 143, +random = 144, +unweight = 145, +vegas_forward = 146, +vegas_inverse = 147, +vegas_histogram = 148, +histogram = 149 diff --git a/madspace/include/madspace/phasespace.hpp b/madspace/include/madspace/phasespace.hpp index 3746db95f2..e2f1ade962 100644 --- a/madspace/include/madspace/phasespace.hpp +++ b/madspace/include/madspace/phasespace.hpp @@ -22,6 +22,7 @@ #include "phasespace/rambo.hpp" #include "phasespace/scale.hpp" #include "phasespace/t_propagator_mapping.hpp" +#include "phasespace/color_ordered_mapping.hpp" #include "phasespace/three_particle.hpp" #include "phasespace/topology.hpp" #include "phasespace/two_particle.hpp" diff --git a/madspace/include/madspace/phasespace/color_ordered_mapping.hpp b/madspace/include/madspace/phasespace/color_ordered_mapping.hpp new file mode 100644 index 0000000000..cd1d70a9c8 --- /dev/null +++ b/madspace/include/madspace/phasespace/color_ordered_mapping.hpp @@ -0,0 +1,55 @@ +#pragma once + +#include + +#include "madspace/phasespace/base.hpp" +#include "madspace/phasespace/invariants.hpp" +#include "madspace/phasespace/three_particle.hpp" +#include "madspace/phasespace/topology.hpp" +#include "madspace/phasespace/two_particle.hpp" + +namespace madspace { + +class ColorOrderedMapping : public Mapping { +public: + // color_order: 0-indexed permutation of {0, ..., n-1} (n = n_out + 2). + // Particles 0 and 1 are the two incoming beams. + ColorOrderedMapping( + const std::vector& color_order, + double t_invariant_power = 0.8, + double s_invariant_power = 0.8 + ); + + std::size_t random_dim() const { return _random_dim; } + +private: + Result build_forward_impl( + FunctionBuilder& fb, + const NamedVector& inputs, + const NamedVector& conditions + ) const override; + Result build_inverse_impl( + FunctionBuilder& fb, + const NamedVector& inputs, + const NamedVector& conditions + ) const override; + + // 0-indexed outgoing-particle indices (values in {0,...,n_out-1}). + // _set1 contains the outgoing particles attached to beam 0's side, + // _set2 those attached to beam 1's side, in peel order. + std::vector _set1; + std::vector _set2; + std::size_t _n_out; + std::size_t _random_dim; + // True iff exactly one of (set1, set2) has size 1 (and the other >= 2). + // In that case the central block is DoubleT instead of 2->2. + bool _use_double_t; + + Invariant _uniform_invariant; + TwoToTwoParticleScattering _com_scattering; + TwoToTwoParticleScattering _lab_scattering; + TwoToThreeParticleScattering _two_to_three; + DoubleT _double_t; +}; + +} // namespace madspace diff --git a/madspace/include/madspace/phasespace/two_particle.hpp b/madspace/include/madspace/phasespace/two_particle.hpp index b038f2c4ad..5faf949942 100644 --- a/madspace/include/madspace/phasespace/two_particle.hpp +++ b/madspace/include/madspace/phasespace/two_particle.hpp @@ -47,4 +47,31 @@ class TwoToTwoParticleScattering : public Mapping { Invariant _invariant; }; +class DoubleT : public Mapping { +public: + DoubleT( + double t1_invariant_power = 0, + double t1_mass = 0, + double t1_width = 0, + double t2_invariant_power = 0, + double t2_mass = 0, + double t2_width = 0 + ); + +private: + Result build_forward_impl( + FunctionBuilder& fb, + const NamedVector& inputs, + const NamedVector& conditions + ) const override; + Result build_inverse_impl( + FunctionBuilder& fb, + const NamedVector& inputs, + const NamedVector& conditions + ) const override; + + Invariant _t1_invariant; + Invariant _t2_invariant; +}; + } // namespace madspace diff --git a/madspace/instruction_set.yaml b/madspace/instruction_set.yaml index ec98a4842a..116d298f26 100644 --- a/madspace/instruction_set.yaml +++ b/madspace/instruction_set.yaml @@ -255,7 +255,7 @@ reduce_sum_vector: batch_reduce_mean: inputs: - - name: in + - name: in1 type: [float] desc: outputs: @@ -267,7 +267,7 @@ batch_reduce_mean: batch_reduce_mean_keepdim: inputs: - - name: in + - name: in1 type: [float] desc: outputs: @@ -964,6 +964,59 @@ two_to_three_particle_scattering_inverse: type: [float] desc: +double_t_scattering: + inputs: + - name: r_phi + type: [float] + desc: + - name: pa + type: [float, 4] + desc: + - name: pb + type: [float, 4] + desc: + - name: t1_abs + type: [float] + desc: + - name: t2_abs + type: [float] + desc: + - name: m1 + type: [float] + desc: + outputs: + - name: p1 + type: [float, 4] + desc: + - name: p2 + type: [float, 4] + desc: + - name: det + type: [float] + desc: + +double_t_scattering_inverse: + inputs: + - name: p1 + type: [float, 4] + desc: + - name: p2 + type: [float, 4] + desc: + - name: pa + type: [float, 4] + desc: + - name: pb + type: [float, 4] + desc: + outputs: + - name: r_phi + type: [float] + desc: + - name: det + type: [float] + desc: + three_body_decay_com: inputs: - name: r_e1 @@ -1189,6 +1242,112 @@ t_inv_value_and_min_max: type: [float] desc: +t1_inv_min_max_doublet: + inputs: + - name: pa + type: [float, 4] + desc: + - name: pb + type: [float, 4] + desc: + - name: m1 + type: [float] + desc: + - name: mir_min + type: [float] + desc: + outputs: + - name: t_min + type: [float] + desc: + - name: t_max + type: [float] + desc: + +t1_inv_value_and_min_max_doublet: + inputs: + - name: pa + type: [float, 4] + desc: + - name: pb + type: [float, 4] + desc: + - name: p1 + type: [float, 4] + desc: + - name: m1 + type: [float] + desc: + - name: mir_min + type: [float] + desc: + outputs: + - name: t_abs + type: [float] + desc: + - name: t_min + type: [float] + desc: + - name: t_max + type: [float] + desc: + +t2_inv_min_max_doublet: + inputs: + - name: pa + type: [float, 4] + desc: + - name: pb + type: [float, 4] + desc: + - name: m1 + type: [float] + desc: + - name: mir_min + type: [float] + desc: + - name: t1_abs + type: [float] + desc: + outputs: + - name: t_min + type: [float] + desc: + - name: t_max + type: [float] + desc: + +t2_inv_value_and_min_max_doublet: + inputs: + - name: pa + type: [float, 4] + desc: + - name: pb + type: [float, 4] + desc: + - name: p1 + type: [float, 4] + desc: + - name: m1 + type: [float] + desc: + - name: mir_min + type: [float] + desc: + - name: t1_abs + type: [float] + desc: + outputs: + - name: t_abs + type: [float] + desc: + - name: t_min + type: [float] + desc: + - name: t_max + type: [float] + desc: + s23_min_max: inputs: - name: pa diff --git a/madspace/madspace/_madspace_py.pyi b/madspace/madspace/_madspace_py.pyi index 963c150df2..ce4e6ad39b 100644 --- a/madspace/madspace/_madspace_py.pyi +++ b/madspace/madspace/_madspace_py.pyi @@ -1,1093 +1,1062 @@ from __future__ import annotations - import collections.abc import typing - -__all__ = [ - "AlphaSGrid", - "BatchSize", - "ChannelWeightNetwork", - "Context", - "CutItem", - "Cuts", - "DataType", - "Decay", - "Device", - "Diagram", - "DifferentialCrossSection", - "DiscreteFlow", - "DiscreteHistogram", - "DiscreteOptimizer", - "DiscreteSampler", - "EnergyScale", - "EventGenerator", - "EventGeneratorConfig", - "EventGeneratorHistogram", - "EventGeneratorStatus", - "EventGeneratorVerbosity", - "FastRamboMapping", - "Flow", - "Function", - "FunctionBuilder", - "FunctionGenerator", - "FunctionRuntime", - "HistItem", - "Instruction", - "InstructionCall", - "Integrand", - "IntegrandProbability", - "Invariant", - "LHECompleter", - "LHEEvent", - "LHEFileWriter", - "LHEHeader", - "LHEMeta", - "LHEParticle", - "LHEProcess", - "LineRef", - "Logger", - "Luminosity", - "MLP", - "Mapping", - "MatrixElement", - "MatrixElementApi", - "MomentumPreprocessing", - "MultiChannelFunction", - "MultiChannelIntegrand", - "MultiChannelMapping", - "Observable", - "ObservableHistograms", - "PartonDensity", - "PdfGrid", - "PhaseSpaceMapping", - "PrettyBox", - "Propagator", - "PropagatorChannelWeights", - "RunningCoupling", - "SubchannelWeights", - "SubprocArgs", - "TPropagatorMapping", - "Tensor", - "ThreeBodyDecay", - "Topology", - "TwoBodyDecay", - "TwoToThreeParticleScattering", - "TwoToTwoParticleScattering", - "Type", - "Unweighter", - "Value", - "VegasGridOptimizer", - "VegasHistogram", - "VegasMapping", - "batch_float", - "batch_float_array", - "batch_four_vec", - "batch_four_vec_array", - "batch_int", - "batch_size", - "batch_sizes", - "cpu_device", - "cuda_device", - "default_context", - "default_cuda_context", - "default_hip_context", - "float", - "format_progress", - "format_si_prefix", - "format_with_error", - "hip_device", - "initialize_vegas_grid", - "int", - "log", - "multichannel_batch_size", - "pretty", - "set_lib_path", - "set_simd_vector_size", - "set_thread_count", - "silent", - "single_float", - "single_int", -] - +__all__: list[str] = ['AdamOptimizer', 'AlphaSGrid', 'BatchSize', 'ChannelEventGenerator', 'ChannelWeightNetwork', 'ColorOrderedMapping', 'Context', 'CutItem', 'Cuts', 'DataType', 'Decay', 'Device', 'Diagram', 'DifferentialCrossSection', 'DiscreteFlow', 'DiscreteHistogram', 'DiscreteOptimizer', 'DiscreteSampler', 'DoubleT', 'EnergyScale', 'EventGenerator', 'FastRamboMapping', 'Flow', 'Function', 'FunctionBuilder', 'FunctionGenerator', 'FunctionRuntime', 'GeneratorConfig', 'GeneratorStatus', 'HistItem', 'Histogram', 'Instruction', 'InstructionCall', 'Integrand', 'IntegrandProbability', 'Invariant', 'LHECompleter', 'LHEEvent', 'LHEFileWriter', 'LHEHeader', 'LHEMeta', 'LHEParticle', 'LHEProcess', 'LineRef', 'Logger', 'Luminosity', 'MLP', 'MadnisConfig', 'MadnisLoss', 'MadnisTraining', 'Mapping', 'MatrixElement', 'MatrixElementApi', 'MomentumPreprocessing', 'MultiChannelFunction', 'MultiChannelIntegrand', 'MultiChannelMapping', 'MultiMadnisTraining', 'NamedTypes', 'NamedValues', 'Observable', 'ObservableHistograms', 'PartonDensity', 'PdfGrid', 'PhaseSpaceMapping', 'PrettyBox', 'Propagator', 'PropagatorChannelWeights', 'RunningCoupling', 'SubchannelWeights', 'SubprocArgs', 'TPropagatorMapping', 'Tensor', 'ThreeBodyDecay', 'Topology', 'TwoBodyDecay', 'TwoToThreeParticleScattering', 'TwoToTwoParticleScattering', 'Type', 'Unweighter', 'Value', 'VegasGridOptimizer', 'VegasHistogram', 'VegasMapping', 'Verbosity', 'batch_float', 'batch_float_array', 'batch_four_vec', 'batch_four_vec_array', 'batch_int', 'batch_size', 'batch_sizes', 'cpu_device', 'cuda_device', 'default_context', 'default_cuda_context', 'default_hip_context', 'float', 'format_progress', 'format_si_prefix', 'format_with_error', 'hip_device', 'initialize_vegas_grid', 'int', 'log', 'multichannel_batch_size', 'pretty', 'set_lib_path', 'set_simd_vector_size', 'silent', 'single_float', 'single_int'] +class AdamOptimizer: + class LRSchedule: + """ + Members: + + none + + cosine + """ + __members__: typing.ClassVar[dict[str, AdamOptimizer.LRSchedule]] # value = {'none': , 'cosine': } + cosine: typing.ClassVar[AdamOptimizer.LRSchedule] # value = + none: typing.ClassVar[AdamOptimizer.LRSchedule] # value = + def __eq__(self, other: typing.Any) -> bool: + ... + def __getstate__(self) -> int: + ... + def __hash__(self) -> int: + ... + def __index__(self) -> int: + ... + @typing.overload + def __init__(self, value: typing.SupportsInt) -> None: + ... + @typing.overload + def __init__(self, name: str) -> None: + ... + def __int__(self) -> int: + ... + def __ne__(self, other: typing.Any) -> bool: + ... + def __repr__(self) -> str: + ... + def __setstate__(self, state: typing.SupportsInt) -> None: + ... + def __str__(self) -> str: + ... + @property + def name(self) -> str: + ... + @property + def value(self) -> int: + ... + cosine: typing.ClassVar[AdamOptimizer.LRSchedule] # value = + none: typing.ClassVar[AdamOptimizer.LRSchedule] # value = + def __init__(self, function: Function, context: Context, learning_rate: typing.SupportsFloat, schedule: AdamOptimizer.LRSchedule = AdamOptimizer.LRSchedule.LRSchedule.none, step_count: typing.SupportsInt = 0, beta1: typing.SupportsFloat = 0.9, beta2: typing.SupportsFloat = 0.999, eps: typing.SupportsFloat = 1e-08) -> None: + ... + def context(self) -> Context: + ... + def input_types(self) -> list[Type]: + ... + def learning_rate(self) -> float: + ... + def step(self, inputs: collections.abc.Sequence[typing.Any]) -> list[Tensor]: + ... class AlphaSGrid: - def __init__(self, file: str) -> None: ... - def coefficients_shape(self, batch_dim: bool = False) -> list[int]: ... - def initialize_globals(self, context: Context, prefix: str = "") -> None: ... - def logq2_shape(self, batch_dim: bool = False) -> list[int]: ... + def __init__(self, file: str) -> None: + ... + def coefficients_shape(self, batch_dim: bool = False) -> list[int]: + ... + def initialize_globals(self, context: Context, prefix: str = '') -> None: + ... + def logq2_shape(self, batch_dim: bool = False) -> list[int]: + ... @property - def logq2(self) -> list[float]: ... + def logq2(self) -> list[float]: + ... @property - def q(self) -> list[float]: ... + def q(self) -> list[float]: + ... @property - def q_count(self) -> int: ... + def q_count(self) -> int: + ... @property - def region_sizes(self) -> list[int]: ... + def region_sizes(self) -> list[int]: + ... @property - def values(self) -> list[float]: ... - + def values(self) -> list[float]: + ... class BatchSize: one: typing.ClassVar[BatchSize] # value = 1 @typing.overload - def __init__(self) -> None: ... + def __init__(self) -> None: + ... @typing.overload - def __init__(self, name: str) -> None: ... - def __repr__(self) -> str: ... - def __str__(self) -> str: ... - + def __init__(self, name: str) -> None: + ... + def __repr__(self) -> str: + ... + def __str__(self) -> str: + ... +class ChannelEventGenerator: + integrand_flags: typing.ClassVar[int] = 1077 + @staticmethod + def load(channel_file: str, contexts: collections.abc.Sequence[Context], event_file: str, weight_file: str, config: GeneratorConfig) -> ChannelEventGenerator: + ... + def __init__(self, contexts: collections.abc.Sequence[Context], integrand: Integrand, event_file: str, weight_file: str, config: GeneratorConfig, subprocess_index: typing.SupportsInt, name: str, histograms: madspace._madspace_py.ObservableHistograms | None) -> None: + ... + def save(self, save: str) -> None: + ... + def status(self) -> GeneratorStatus: + ... class ChannelWeightNetwork(FunctionGenerator): - def __init__( - self, - channel_count: typing.SupportsInt, - particle_count: typing.SupportsInt, - hidden_dim: typing.SupportsInt = 32, - layers: typing.SupportsInt = 3, - activation: MLP.Activation = MLP.Activation.Activation.leaky_relu, - prefix: str = "", - ) -> None: ... - def initialize_globals(self, context: Context) -> None: ... - def mask_name(self) -> str: ... - def mlp(self) -> MLP: ... - def preprocessing(self) -> MomentumPreprocessing: ... - + def __init__(self, channel_count: typing.SupportsInt, particle_count: typing.SupportsInt, hidden_dim: typing.SupportsInt = 32, layers: typing.SupportsInt = 3, activation: MLP.Activation = MLP.Activation.Activation.leaky_relu, prefix: str = '', include_preprocessing: bool = True) -> None: + ... + def initialize_globals(self, context: Context) -> None: + ... + def mask_name(self) -> str: + ... + def mlp(self) -> MLP: + ... + def preprocessing(self) -> MomentumPreprocessing: + ... +class ColorOrderedMapping(Mapping): + def __init__(self, colour_order: collections.abc.Sequence[typing.SupportsInt], t_invariant_power: typing.SupportsFloat = 0.0, s_invariant_power: typing.SupportsFloat = 0.0) -> None: + ... + def random_dim(self) -> int: + ... class Context: @typing.overload - def __init__(self) -> None: ... + def __init__(self, thread_count: typing.SupportsInt = -1) -> None: + ... @typing.overload - def __init__(self, device: Device) -> None: ... - def define_global( - self, - name: str, - dtype: DataType, - shape: collections.abc.Sequence[typing.SupportsInt], - requires_grad: bool = False, - ) -> Tensor: ... - def device(self) -> Device: ... - def get_global(self, name: str) -> Tensor: ... - def global_exists(self, name: str) -> bool: ... - def global_requires_grad(self, name: str) -> bool: ... - def load(self, file: str) -> None: ... - def load_matrix_element(self, file: str, param_card: str) -> MatrixElementApi: ... - def matrix_element(self, index: typing.SupportsInt) -> MatrixElementApi: ... - def save(self, file: str) -> None: ... - + def __init__(self, device: Device, thread_count: typing.SupportsInt = -1) -> None: + ... + def copy_globals_from(self, context: Context) -> None: + ... + def define_global(self, name: str, dtype: DataType, shape: collections.abc.Sequence[typing.SupportsInt], requires_grad: bool = False) -> Tensor: + ... + def delete_global(self, name: str) -> None: + ... + def device(self) -> Device: + ... + def get_global(self, name: str) -> Tensor: + ... + def global_exists(self, name: str) -> bool: + ... + def global_names(self) -> list[str]: + ... + def global_requires_grad(self, name: str) -> bool: + ... + def load_globals(self, dir: str) -> None: + ... + def load_matrix_element(self, file: str, param_card: str) -> MatrixElementApi: + ... + def matrix_element(self, index: typing.SupportsInt) -> MatrixElementApi: + ... + def save_globals(self, dir: str) -> None: + ... class CutItem: - def __init__( - self, - observable: Observable, - min: typing.SupportsFloat = ..., - max: typing.SupportsFloat = ..., - mode: Cuts.CutMode = Cuts.CutMode.CutMode.all, - ) -> None: ... + def __init__(self, observable: Observable, min: typing.SupportsFloat = ..., max: typing.SupportsFloat = ..., mode: Cuts.CutMode = Cuts.CutMode.CutMode.all) -> None: + ... @property - def max(self) -> float: ... + def max(self) -> float: + ... @property - def min(self) -> float: ... + def min(self) -> float: + ... @property - def mode(self) -> Cuts.CutMode: ... + def mode(self) -> Cuts.CutMode: + ... @property - def observable(self) -> Observable: ... - + def observable(self) -> Observable: + ... class Cuts(FunctionGenerator): class CutMode: """ Members: - + any - + all """ - - __members__: typing.ClassVar[ - dict[str, Cuts.CutMode] - ] # value = {'any': , 'all': } + __members__: typing.ClassVar[dict[str, Cuts.CutMode]] # value = {'any': , 'all': } all: typing.ClassVar[Cuts.CutMode] # value = any: typing.ClassVar[Cuts.CutMode] # value = - def __eq__(self, other: typing.Any) -> bool: ... - def __getstate__(self) -> int: ... - def __hash__(self) -> int: ... - def __index__(self) -> int: ... + def __eq__(self, other: typing.Any) -> bool: + ... + def __getstate__(self) -> int: + ... + def __hash__(self) -> int: + ... + def __index__(self) -> int: + ... @typing.overload - def __init__(self, value: typing.SupportsInt) -> None: ... + def __init__(self, value: typing.SupportsInt) -> None: + ... @typing.overload - def __init__(self, name: str) -> None: ... - def __int__(self) -> int: ... - def __ne__(self, other: typing.Any) -> bool: ... - def __repr__(self) -> str: ... - def __setstate__(self, state: typing.SupportsInt) -> None: ... - def __str__(self) -> str: ... + def __init__(self, name: str) -> None: + ... + def __int__(self) -> int: + ... + def __ne__(self, other: typing.Any) -> bool: + ... + def __repr__(self) -> str: + ... + def __setstate__(self, state: typing.SupportsInt) -> None: + ... + def __str__(self) -> str: + ... @property - def name(self) -> str: ... + def name(self) -> str: + ... @property - def value(self) -> int: ... - + def value(self) -> int: + ... all: typing.ClassVar[Cuts.CutMode] # value = any: typing.ClassVar[Cuts.CutMode] # value = @typing.overload - def __init__(self, cut_data: collections.abc.Sequence[CutItem]) -> None: ... + def __init__(self, cut_data: collections.abc.Sequence[CutItem]) -> None: + ... @typing.overload - def __init__(self, particle_count: typing.SupportsInt) -> None: ... - def eta_max(self) -> list[float]: ... - def pt_min(self) -> list[float]: ... - def sqrt_s_min(self) -> float: ... - + def __init__(self, particle_count: typing.SupportsInt) -> None: + ... + def eta_max(self) -> list[float]: + ... + def pt_min(self) -> list[float]: + ... + def sqrt_s_min(self) -> float: + ... class DataType: """ Members: - + int - + float - + batch_sizes """ - - __members__: typing.ClassVar[ - dict[str, DataType] - ] # value = {'int': , 'float': , 'batch_sizes': } + __members__: typing.ClassVar[dict[str, DataType]] # value = {'int': , 'float': , 'batch_sizes': } batch_sizes: typing.ClassVar[DataType] # value = float: typing.ClassVar[DataType] # value = int: typing.ClassVar[DataType] # value = - def __eq__(self, other: typing.Any) -> bool: ... - def __getstate__(self) -> int: ... - def __hash__(self) -> int: ... - def __index__(self) -> int: ... + def __eq__(self, other: typing.Any) -> bool: + ... + def __getstate__(self) -> int: + ... + def __hash__(self) -> int: + ... + def __index__(self) -> int: + ... @typing.overload - def __init__(self, value: typing.SupportsInt) -> None: ... + def __init__(self, value: typing.SupportsInt) -> None: + ... @typing.overload - def __init__(self, name: str) -> None: ... - def __int__(self) -> int: ... - def __ne__(self, other: typing.Any) -> bool: ... - def __repr__(self) -> str: ... - def __setstate__(self, state: typing.SupportsInt) -> None: ... - def __str__(self) -> str: ... - @property - def name(self) -> str: ... - @property - def value(self) -> int: ... - + def __init__(self, name: str) -> None: + ... + def __int__(self) -> int: + ... + def __ne__(self, other: typing.Any) -> bool: + ... + def __repr__(self) -> str: + ... + def __setstate__(self, state: typing.SupportsInt) -> None: + ... + def __str__(self) -> str: + ... + @property + def name(self) -> str: + ... + @property + def value(self) -> int: + ... class Decay: @property - def child_indices(self) -> list[int]: ... + def child_indices(self) -> list[int]: + ... @property - def e_max(self) -> float: ... + def e_max(self) -> float: + ... @property - def e_min(self) -> float: ... + def e_min(self) -> float: + ... @property - def index(self) -> int: ... + def index(self) -> int: + ... @property - def mass(self) -> float: ... + def mass(self) -> float: + ... @property - def on_shell(self) -> bool: ... + def on_shell(self) -> bool: + ... @property - def parent_index(self) -> int: ... + def parent_index(self) -> int: + ... @property - def pdg_id(self) -> int: ... + def pdg_id(self) -> int: + ... @property - def width(self) -> float: ... - + def width(self) -> float: + ... class Device: pass - class Diagram: - def __init__( - self, - incoming_masses: collections.abc.Sequence[typing.SupportsFloat], - outgoing_masses: collections.abc.Sequence[typing.SupportsFloat], - propagators: collections.abc.Sequence[Propagator], - vertices: collections.abc.Sequence[collections.abc.Sequence[LineRef]], - ) -> None: ... + def __init__(self, incoming_masses: collections.abc.Sequence[typing.SupportsFloat], outgoing_masses: collections.abc.Sequence[typing.SupportsFloat], propagators: collections.abc.Sequence[Propagator], vertices: collections.abc.Sequence[collections.abc.Sequence[LineRef]]) -> None: + ... @property - def incoming_masses(self) -> list[float]: ... + def incoming_masses(self) -> list[float]: + ... @property - def incoming_vertices(self) -> typing.Annotated[list[int], "FixedSize(2)"]: ... + def incoming_vertices(self) -> typing.Annotated[list[int], "FixedSize(2)"]: + ... @property - def outgoing_masses(self) -> list[float]: ... + def outgoing_masses(self) -> list[float]: + ... @property - def outgoing_vertices(self) -> list[int]: ... + def outgoing_vertices(self) -> list[int]: + ... @property - def propagator_vertices(self) -> list[list[int]]: ... + def propagator_vertices(self) -> list[list[int]]: + ... @property - def propagators(self) -> list[Propagator]: ... + def propagators(self) -> list[Propagator]: + ... @property - def vertices(self) -> list[list[LineRef]]: ... - + def vertices(self) -> list[list[LineRef]]: + ... class DifferentialCrossSection(FunctionGenerator): - def __init__( - self, - matrix_element: MatrixElement, - cm_energy: typing.SupportsFloat, - running_coupling: RunningCoupling, - energy_scale: EnergyScale, - pid_options: collections.abc.Sequence[ - collections.abc.Sequence[typing.SupportsInt] - ] = [], - has_pdf1: bool = False, - has_pdf2: bool = False, - pdf_grid1: madspace._madspace_py.PdfGrid | None = None, - pdf_grid2: madspace._madspace_py.PdfGrid | None = None, - has_mirror: bool = False, - input_momentum_fraction: bool = True, - ) -> None: ... - def has_mirror(self) -> bool: ... - def matrix_element(self) -> MatrixElement: ... - def pid_options(self) -> list[list[int]]: ... - + def __init__(self, matrix_element: MatrixElement, cm_energy: typing.SupportsFloat, running_coupling: RunningCoupling, energy_scale: EnergyScale, pid_options: collections.abc.Sequence[collections.abc.Sequence[typing.SupportsInt]] = [], has_pdf1: bool = False, has_pdf2: bool = False, pdf_grid1: madspace._madspace_py.PdfGrid | None = None, pdf_grid2: madspace._madspace_py.PdfGrid | None = None, has_mirror: bool = False, input_momentum_fraction: bool = True) -> None: + ... + def has_mirror(self) -> bool: + ... + def matrix_element(self) -> MatrixElement: + ... + def pid_options(self) -> list[list[int]]: + ... class DiscreteFlow(Mapping): - def __init__( - self, - option_counts: collections.abc.Sequence[typing.SupportsInt], - prefix: str = "", - dims_with_prior: collections.abc.Sequence[typing.SupportsInt] = [], - condition_dim: typing.SupportsInt = 0, - subnet_hidden_dim: typing.SupportsInt = 32, - subnet_layers: typing.SupportsInt = 3, - subnet_activation: MLP.Activation = MLP.Activation.Activation.leaky_relu, - ) -> None: ... - def condition_dim(self) -> int: ... - def initialize_globals(self, context: Context) -> None: ... - def option_counts(self) -> list[int]: ... - + def __init__(self, option_counts: collections.abc.Sequence[typing.SupportsInt], prefix: str = '', dims_with_prior: collections.abc.Sequence[typing.SupportsInt] = [], condition_dim: typing.SupportsInt = 0, subnet_hidden_dim: typing.SupportsInt = 32, subnet_layers: typing.SupportsInt = 3, subnet_activation: MLP.Activation = MLP.Activation.Activation.leaky_relu) -> None: + ... + def condition_dim(self) -> int: + ... + def initialize_globals(self, context: Context) -> None: + ... + def option_counts(self) -> list[int]: + ... class DiscreteHistogram(FunctionGenerator): - def __init__( - self, option_counts: collections.abc.Sequence[typing.SupportsInt] - ) -> None: ... - + def __init__(self, option_counts: collections.abc.Sequence[typing.SupportsInt]) -> None: + ... class DiscreteOptimizer: - def __init__( - self, context: Context, prob_names: collections.abc.Sequence[str] - ) -> None: ... - def add_data( - self, values_and_counts: collections.abc.Sequence[typing.Any] - ) -> None: ... - def optimize(self) -> None: ... - + def __init__(self, contexts: collections.abc.Sequence[Context], prob_names: collections.abc.Sequence[str]) -> None: + ... + def add_data(self, values_and_counts: collections.abc.Sequence[typing.Any]) -> None: + ... + def optimize(self) -> None: + ... class DiscreteSampler(Mapping): - def __init__( - self, - option_counts: collections.abc.Sequence[typing.SupportsInt], - prefix: str = "", - dims_with_prior: collections.abc.Sequence[typing.SupportsInt] = [], - ) -> None: ... - def initialize_globals(self, context: Context) -> None: ... - def option_counts(self) -> list[int]: ... - def prob_names(self) -> list[str]: ... - + def __init__(self, option_counts: collections.abc.Sequence[typing.SupportsInt], prefix: str = '', dims_with_prior: collections.abc.Sequence[typing.SupportsInt] = []) -> None: + ... + def initialize_globals(self, context: Context) -> None: + ... + def option_counts(self) -> list[int]: + ... + def prob_names(self) -> list[str]: + ... +class DoubleT(Mapping): + def __init__(self, t1_invariant_power: typing.SupportsFloat = 0.0, t1_mass: typing.SupportsFloat = 0.0, t1_width: typing.SupportsFloat = 0.0, t2_invariant_power: typing.SupportsFloat = 0.0, t2_mass: typing.SupportsFloat = 0.0, t2_width: typing.SupportsFloat = 0.0) -> None: + ... class EnergyScale(FunctionGenerator): class DynamicalScaleType: """ Members: - + transverse_energy - + transverse_mass - + half_transverse_mass - + partonic_energy """ - - __members__: typing.ClassVar[ - dict[str, EnergyScale.DynamicalScaleType] - ] # value = {'transverse_energy': , 'transverse_mass': , 'half_transverse_mass': , 'partonic_energy': } - half_transverse_mass: typing.ClassVar[ - EnergyScale.DynamicalScaleType - ] # value = - partonic_energy: typing.ClassVar[ - EnergyScale.DynamicalScaleType - ] # value = - transverse_energy: typing.ClassVar[ - EnergyScale.DynamicalScaleType - ] # value = - transverse_mass: typing.ClassVar[ - EnergyScale.DynamicalScaleType - ] # value = - def __eq__(self, other: typing.Any) -> bool: ... - def __getstate__(self) -> int: ... - def __hash__(self) -> int: ... - def __index__(self) -> int: ... + __members__: typing.ClassVar[dict[str, EnergyScale.DynamicalScaleType]] # value = {'transverse_energy': , 'transverse_mass': , 'half_transverse_mass': , 'partonic_energy': } + half_transverse_mass: typing.ClassVar[EnergyScale.DynamicalScaleType] # value = + partonic_energy: typing.ClassVar[EnergyScale.DynamicalScaleType] # value = + transverse_energy: typing.ClassVar[EnergyScale.DynamicalScaleType] # value = + transverse_mass: typing.ClassVar[EnergyScale.DynamicalScaleType] # value = + def __eq__(self, other: typing.Any) -> bool: + ... + def __getstate__(self) -> int: + ... + def __hash__(self) -> int: + ... + def __index__(self) -> int: + ... @typing.overload - def __init__(self, value: typing.SupportsInt) -> None: ... + def __init__(self, value: typing.SupportsInt) -> None: + ... @typing.overload - def __init__(self, name: str) -> None: ... - def __int__(self) -> int: ... - def __ne__(self, other: typing.Any) -> bool: ... - def __repr__(self) -> str: ... - def __setstate__(self, state: typing.SupportsInt) -> None: ... - def __str__(self) -> str: ... + def __init__(self, name: str) -> None: + ... + def __int__(self) -> int: + ... + def __ne__(self, other: typing.Any) -> bool: + ... + def __repr__(self) -> str: + ... + def __setstate__(self, state: typing.SupportsInt) -> None: + ... + def __str__(self) -> str: + ... @property - def name(self) -> str: ... + def name(self) -> str: + ... @property - def value(self) -> int: ... - - half_transverse_mass: typing.ClassVar[ - EnergyScale.DynamicalScaleType - ] # value = - partonic_energy: typing.ClassVar[ - EnergyScale.DynamicalScaleType - ] # value = - transverse_energy: typing.ClassVar[ - EnergyScale.DynamicalScaleType - ] # value = - transverse_mass: typing.ClassVar[ - EnergyScale.DynamicalScaleType - ] # value = + def value(self) -> int: + ... + half_transverse_mass: typing.ClassVar[EnergyScale.DynamicalScaleType] # value = + partonic_energy: typing.ClassVar[EnergyScale.DynamicalScaleType] # value = + transverse_energy: typing.ClassVar[EnergyScale.DynamicalScaleType] # value = + transverse_mass: typing.ClassVar[EnergyScale.DynamicalScaleType] # value = @typing.overload - def __init__(self, particle_count: typing.SupportsInt) -> None: ... + def __init__(self, particle_count: typing.SupportsInt) -> None: + ... @typing.overload - def __init__( - self, particle_count: typing.SupportsInt, type: EnergyScale.DynamicalScaleType - ) -> None: ... + def __init__(self, particle_count: typing.SupportsInt, type: EnergyScale.DynamicalScaleType) -> None: + ... @typing.overload - def __init__( - self, particle_count: typing.SupportsInt, fixed_scale: typing.SupportsFloat - ) -> None: ... + def __init__(self, particle_count: typing.SupportsInt, fixed_scale: typing.SupportsFloat) -> None: + ... @typing.overload - def __init__( - self, - particle_count: typing.SupportsInt, - dynamical_scale_type: EnergyScale.DynamicalScaleType, - ren_scale_fixed: bool, - fact_scale_fixed: bool, - ren_scale: typing.SupportsFloat, - fact_scale1: typing.SupportsFloat, - fact_scale2: typing.SupportsFloat, - ) -> None: ... - + def __init__(self, particle_count: typing.SupportsInt, dynamical_scale_type: EnergyScale.DynamicalScaleType, ren_scale_fixed: bool, fact_scale_fixed: bool, ren_scale: typing.SupportsFloat, fact_scale1: typing.SupportsFloat, fact_scale2: typing.SupportsFloat) -> None: + ... class EventGenerator: - default_config: typing.ClassVar[ - EventGeneratorConfig - ] # value = - integrand_flags: typing.ClassVar[int] = 1077 - def __init__( - self, - context: Context, - channels: collections.abc.Sequence[Integrand], - temp_file_prefix: str, - status_file: str = "", - config: EventGeneratorConfig = ..., - channel_subprocesses: collections.abc.Sequence[typing.SupportsInt] = [], - channel_names: collections.abc.Sequence[str] = [], - channel_histograms: collections.abc.Sequence[ObservableHistograms] = [], - ) -> None: ... - def channel_status(self) -> list[EventGeneratorStatus]: ... - def combine_to_compact_npy(self, file_name: str) -> None: ... - def combine_to_lhe(self, file_name: str, lhe_completer: ...) -> None: ... - def combine_to_lhe_npy(self, file_name: str, lhe_completer: ...) -> None: ... - def generate(self) -> None: ... - def histograms(self) -> list[EventGeneratorHistogram]: ... - def status(self) -> EventGeneratorStatus: ... - def survey(self) -> None: ... - -class EventGeneratorConfig: - verbosity: EventGeneratorVerbosity - def __init__(self) -> None: ... - @property - def batch_size(self) -> int: ... - @batch_size.setter - def batch_size(self, arg0: typing.SupportsInt) -> None: ... - @property - def freeze_max_weight_after(self) -> int: ... + default_config: typing.ClassVar[GeneratorConfig] # value = + def __init__(self, contexts: collections.abc.Sequence[Context], channels: collections.abc.Sequence[ChannelEventGenerator], status_file: str = '', config: GeneratorConfig = ...) -> None: + ... + def channel_status(self) -> list[GeneratorStatus]: + ... + def channels(self) -> list[ChannelEventGenerator]: + ... + def combine_to_compact_npy(self, file_name: str) -> None: + ... + def combine_to_lhe(self, file_name: str, lhe_completer: LHECompleter) -> None: + ... + def combine_to_lhe_npy(self, file_name: str, lhe_completer: LHECompleter) -> None: + ... + def generate(self) -> None: + ... + def histograms(self) -> list[Histogram]: + ... + def status(self) -> GeneratorStatus: + ... + def survey(self) -> None: + ... + def used_globals(self) -> set[str]: + ... +class FastRamboMapping(Mapping): + def __init__(self, n_particles: typing.SupportsInt, massless: bool) -> None: + ... +class Flow(Mapping): + def __init__(self, input_dim: typing.SupportsInt, condition_dim: typing.SupportsInt = 0, prefix: str = '', bin_count: typing.SupportsInt = 10, subnet_hidden_dim: typing.SupportsInt = 32, subnet_layers: typing.SupportsInt = 3, subnet_activation: MLP.Activation = MLP.Activation.Activation.leaky_relu, invert_spline: bool = True) -> None: + ... + def condition_dim(self) -> int: + ... + def initialize_from_vegas(self, context: Context, grid_name: str) -> None: + ... + def initialize_globals(self, context: Context) -> None: + ... + def input_dim(self) -> int: + ... +class Function: + @staticmethod + def load(file: str) -> Function: + ... + def __call__(self, *args, **kwargs): + ... + def __repr__(self) -> str: + ... + def __str__(self) -> str: + ... + def save(self, file: str) -> None: + ... + @property + def globals(self) -> list[tuple[str, Value]]: + ... + @property + def inputs(self) -> NamedValues: + ... + @property + def instructions(self) -> list[InstructionCall]: + ... + @property + def locals(self) -> list[Value]: + ... + @property + def outputs(self) -> NamedValues: + ... +class FunctionBuilder: + def __init__(self, input_types: NamedTypes, output_types: NamedTypes) -> None: + ... + def accept_norm(self, accepted_batch: Value, full_batch: Value) -> Value: + ... + def add(self, in1: Value, in2: Value) -> Value: + ... + def add_int(self, in1: Value, in2: Value) -> Value: + ... + def apply_subchannel_weights(self, channel_weights_in: Value, subchannel_weights: Value, channel_indices: Value, subchannel_indices: Value) -> Value: + ... + def argsort(self, input: Value) -> Value: + ... + def batch_cat(self, args: collections.abc.Sequence[Value]) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... + def batch_gather(self, indices: Value, values: Value) -> Value: + ... + def batch_reduce_mean(self, in1: Value) -> Value: + ... + def batch_reduce_mean_keepdim(self, in1: Value) -> Value: + ... + def batch_scatter(self, indices: Value, target: Value, source: Value) -> Value: + ... + def batch_size(self, args: collections.abc.Sequence[Value]) -> Value: + ... + def batch_split(self, in1: Value, counts: Value) -> list[Value]: + ... + def boost_beam(self, p1: Value, x1: Value, x2: Value) -> Value: + ... + def boost_beam_inverse(self, p1: Value, x1: Value, x2: Value) -> Value: + ... + def breit_wigner_invariant(self, r: Value, mass: Value, width: Value, s_min: Value, s_max: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... + def breit_wigner_invariant_inverse(self, s: Value, mass: Value, width: Value, s_min: Value, s_max: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... + def cat(self, args: collections.abc.Sequence[Value]) -> Value: + ... + def chili_forward(self, r: Value, e_cm: Value, m_out: Value, pt_min: Value, y_max: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... + def chili_inverse(self, p_ext: Value, e_cm: Value, m_out: Value, pt_min: Value, y_max: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... + def collect_channel_weights(self, amp2: Value, channel_indices: Value, channel_count: Value) -> Value: + ... + def com_p_in(self, e_cm: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... + def current_stream(self) -> int: + ... + def cut_all(self, obs: Value, min: Value, max: Value) -> Value: + ... + def cut_any(self, obs: Value, min: Value, max: Value) -> Value: + ... + def cut_one(self, obs: Value, min: Value, max: Value) -> Value: + ... + def cut_unphysical(self, w_in: Value, p: Value, x1: Value, x2: Value) -> Value: + ... + def diff_cross_section(self, x1: Value, x2: Value, pdf1: Value, pdf2: Value, matrix_element: Value, e_cm2: Value) -> Value: + ... + def discrete_histogram(self, input: Value, weights: Value, option_count: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... + def div(self, in1: Value, in2: Value) -> Value: + ... + def double_t_scattering(self, r_phi: Value, pa: Value, pb: Value, t1_abs: Value, t2_abs: Value, m1: Value) -> typing.Annotated[list[Value], "FixedSize(3)"]: + ... + def double_t_scattering_inverse(self, p1: Value, p2: Value, pa: Value, pb: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... + def elu(self, in1: Value) -> Value: + ... + def fast_rambo_massive(self, r: Value, e_cm: Value, masses: Value, p0: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... + def fast_rambo_massive_com(self, r: Value, e_cm: Value, masses: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... + def fast_rambo_massive_inverse(self, p_out: Value, e_cm: Value, masses: Value) -> typing.Annotated[list[Value], "FixedSize(3)"]: + ... + def fast_rambo_massless(self, r: Value, e_cm: Value, p0: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... + def fast_rambo_massless_com(self, r: Value, e_cm: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... + def fast_rambo_massless_inverse(self, p_out: Value, e_cm: Value) -> typing.Annotated[list[Value], "FixedSize(3)"]: + ... + def full(self, args: collections.abc.Sequence[Value]) -> Value: + ... + def function(self) -> Function: + ... + def gather(self, index: Value, choices: Value) -> Value: + ... + def gather_int(self, index: Value, choices: Value) -> Value: + ... + def gelu(self, in1: Value) -> Value: + ... + def get_global(self, name: str, dtype: DataType, shape: collections.abc.Sequence[typing.SupportsInt]) -> Value: + ... + def histogram(self, input: Value, weights: Value, min: Value, max: Value, bin_count: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... + def input(self, index: typing.SupportsInt) -> Value: + ... + def input_range(self, start_index: typing.SupportsInt, end_index: typing.SupportsInt) -> list[Value]: + ... + def interpolate_alpha_s(self, q2: Value, grid_logq2: Value, grid_coeffs: Value) -> Value: + ... + def interpolate_pdf(self, x: Value, q2: Value, pid_indices: Value, grid_logx: Value, grid_logq2: Value, grid_coeffs: Value) -> Value: + ... + def invariants_from_momenta(self, p_ext: Value, factors: Value) -> Value: + ... + def leaky_relu(self, in1: Value) -> Value: + ... + def madnis_abs_weight(self, f: Value, q: Value) -> Value: + ... + def madnis_multi_channel_variance(self, vars: Value, abs_means: Value) -> Value: + ... + def madnis_single_channel_variance(self, var: Value, abs_mean: Value) -> Value: + ... + def madnis_softclip(self, f: Value, q: Value, norm: Value, threshold: Value) -> Value: + ... + def madnis_variance(self, f: Value, g: Value, q: Value, mean: Value) -> Value: + ... + def matmul(self, x: Value, weight: Value, bias: Value) -> Value: + ... + def matrix_element(self, args: collections.abc.Sequence[Value]) -> list[Value]: + ... + def max(self, in1: Value, in2: Value) -> Value: + ... + def min(self, in1: Value, in2: Value) -> Value: + ... + def mirror_momenta(self, p_ext: Value, mirror: Value) -> Value: + ... + def momenta_to_x1x2(self, p_ext: Value, e_cm: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... + def mul(self, in1: Value, in2: Value) -> Value: + ... + def nonzero(self, input: Value) -> Value: + ... + def obs_delta_eta(self, p1: Value, p2: Value) -> Value: + ... + def obs_delta_phi(self, p1: Value, p2: Value) -> Value: + ... + def obs_delta_r(self, p1: Value, p2: Value) -> Value: + ... + def obs_e(self, p: Value) -> Value: + ... + def obs_eta(self, p: Value) -> Value: + ... + def obs_eta_abs(self, p: Value) -> Value: + ... + def obs_mass(self, p: Value) -> Value: + ... + def obs_p_mag(self, p: Value) -> Value: + ... + def obs_phi(self, p: Value) -> Value: + ... + def obs_pt(self, p: Value) -> Value: + ... + def obs_px(self, p: Value) -> Value: + ... + def obs_py(self, p: Value) -> Value: + ... + def obs_pz(self, p: Value) -> Value: + ... + def obs_sqrt_s(self, p_ext: Value) -> Value: + ... + def obs_theta(self, p: Value) -> Value: + ... + def obs_y(self, p: Value) -> Value: + ... + def obs_y_abs(self, p: Value) -> Value: + ... + def offset_indices(self, batch_sizes_offset: Value, batch_sizes_out: Value) -> Value: + ... + def one_hot(self, index: Value, option_count: Value) -> Value: + ... + def output(self, index: typing.SupportsInt, value: Value) -> None: + ... + def output_range(self, start_index: typing.SupportsInt, values: collections.abc.Sequence[Value]) -> None: + ... + def permute_momenta(self, momenta: Value, permutations: Value, index: Value) -> Value: + ... + def pop(self, in1: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... + def product(self, values: collections.abc.Sequence[Value]) -> Value: + ... + def pt_eta_phi_x(self, p_ext: Value, x1: Value, x2: Value) -> Value: + ... + def r_to_x1x2(self, r: Value, s_hat: Value, s_lab: Value) -> typing.Annotated[list[Value], "FixedSize(3)"]: + ... + def random(self, batch_size: Value, count: Value) -> Value: + ... + def reduce_product(self, in1: Value) -> Value: + ... + def reduce_sum(self, in1: Value) -> Value: + ... + def reduce_sum_vector(self, in1: Value) -> Value: + ... + def relu(self, in1: Value) -> Value: + ... + def rqs_find_bin(self, input: Value, in_sizes: Value, out_sizes: Value, derivatives: Value) -> Value: + ... + def rqs_forward(self, input: Value, condition: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... + def rqs_inverse(self, input: Value, condition: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... + def rqs_reshape(self, input: Value, bin_count: Value) -> typing.Annotated[list[Value], "FixedSize(3)"]: + ... + def s23_min_max(self, pa: Value, pb: Value, p3: Value, t1_abs: Value, m1: Value, m2: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... + def s23_value_and_min_max(self, pa: Value, pb: Value, p3: Value, t1_abs: Value, p1: Value, p2: Value) -> typing.Annotated[list[Value], "FixedSize(3)"]: + ... + def sample_discrete(self, r: Value, option_count: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... + def sample_discrete_inverse(self, index: Value, option_count: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... + def sample_discrete_probs(self, r: Value, probs: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... + def sample_discrete_probs_inverse(self, index: Value, probs: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... + def scale_half_transverse_mass(self, momenta: Value) -> Value: + ... + def scale_partonic_energy(self, momenta: Value) -> Value: + ... + def scale_transverse_energy(self, momenta: Value) -> Value: + ... + def scale_transverse_mass(self, momenta: Value) -> Value: + ... + def sde2_channel_weights(self, invariants: Value, masses: Value, widths: Value, indices: Value) -> Value: + ... + def select(self, input: Value, indices: Value) -> Value: + ... + def select_int(self, input: Value, indices: Value) -> Value: + ... + def select_vector(self, input: Value, indices: Value) -> Value: + ... + def set_current_stream(self, arg0: typing.SupportsInt) -> None: + ... + def sigmoid(self, in1: Value) -> Value: + ... + def softmax(self, input: Value) -> Value: + ... + def softmax_prior(self, input: Value, prior: Value) -> Value: + ... + def softplus(self, in1: Value) -> Value: + ... + def sqrt(self, in1: Value) -> Value: + ... + def square(self, in1: Value) -> Value: + ... + def squeeze(self, input: Value) -> Value: + ... + def stable_invariant(self, r: Value, mass: Value, s_min: Value, s_max: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... + def stable_invariant_inverse(self, s: Value, mass: Value, s_min: Value, s_max: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... + def stable_invariant_nu(self, r: Value, mass: Value, nu: Value, s_min: Value, s_max: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... + def stable_invariant_nu_inverse(self, s: Value, mass: Value, nu: Value, s_min: Value, s_max: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... + def stack(self, args: collections.abc.Sequence[Value]) -> Value: + ... + def stack_sizes(self, args: collections.abc.Sequence[Value]) -> Value: + ... + def sub(self, in1: Value, in2: Value) -> Value: + ... + def subchannel_weights(self, invariants: Value, masses: Value, widths: Value, indices: Value, on_shell: Value, group_sizes: Value) -> Value: + ... + def t1_inv_min_max_doublet(self, pa: Value, pb: Value, m1: Value, mir_min: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... + def t1_inv_value_and_min_max_doublet(self, pa: Value, pb: Value, p1: Value, m1: Value, mir_min: Value) -> typing.Annotated[list[Value], "FixedSize(3)"]: + ... + def t2_inv_min_max_doublet(self, pa: Value, pb: Value, m1: Value, mir_min: Value, t1_abs: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... + def t2_inv_value_and_min_max_doublet(self, pa: Value, pb: Value, p1: Value, m1: Value, mir_min: Value, t1_abs: Value) -> typing.Annotated[list[Value], "FixedSize(3)"]: + ... + def t_inv_min_max(self, pa: Value, pb: Value, m1: Value, m2: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... + def t_inv_value_and_min_max(self, pa: Value, pb: Value, p1: Value, p2: Value) -> typing.Annotated[list[Value], "FixedSize(3)"]: + ... + def three_body_decay(self, r_e1: Value, r_e2: Value, r_phi: Value, r_cos_theta: Value, r_beta: Value, m0: Value, m1: Value, m2: Value, m3: Value, p0: Value) -> typing.Annotated[list[Value], "FixedSize(4)"]: + ... + def three_body_decay_com(self, r_e1: Value, r_e2: Value, r_phi: Value, r_cos_theta: Value, r_beta: Value, m0: Value, m1: Value, m2: Value, m3: Value) -> typing.Annotated[list[Value], "FixedSize(4)"]: + ... + def three_body_decay_com_inverse(self, p1: Value, p2: Value, p3: Value) -> typing.Annotated[list[Value], "FixedSize(10)"]: + ... + def three_body_decay_inverse(self, p1: Value, p2: Value, p3: Value) -> typing.Annotated[list[Value], "FixedSize(11)"]: + ... + def two_body_decay(self, r_phi: Value, r_cos_theta: Value, m0: Value, m1: Value, m2: Value, p0: Value) -> typing.Annotated[list[Value], "FixedSize(3)"]: + ... + def two_body_decay_com(self, r_phi: Value, r_cos_theta: Value, m0: Value, m1: Value, m2: Value) -> typing.Annotated[list[Value], "FixedSize(3)"]: + ... + def two_body_decay_com_inverse(self, p1: Value, p2: Value) -> typing.Annotated[list[Value], "FixedSize(6)"]: + ... + def two_body_decay_inverse(self, p1: Value, p2: Value) -> typing.Annotated[list[Value], "FixedSize(7)"]: + ... + def two_to_three_particle_scattering(self, phi_choice: Value, pa: Value, pb: Value, p3: Value, s23: Value, t1_abs: Value, m1: Value, m2: Value) -> typing.Annotated[list[Value], "FixedSize(3)"]: + ... + def two_to_three_particle_scattering_inverse(self, p1: Value, p2: Value, p3: Value, pa: Value, pb: Value, t1_abs: Value, s23: Value) -> typing.Annotated[list[Value], "FixedSize(4)"]: + ... + def two_to_two_particle_scattering(self, r_phi: Value, pa: Value, pb: Value, t_abs: Value, m1: Value, m2: Value) -> typing.Annotated[list[Value], "FixedSize(3)"]: + ... + def two_to_two_particle_scattering_com(self, r_phi: Value, pa: Value, pb: Value, t_abs: Value, m1: Value, m2: Value) -> typing.Annotated[list[Value], "FixedSize(3)"]: + ... + def two_to_two_particle_scattering_com_inverse(self, p1: Value, p2: Value, pa: Value, pb: Value) -> typing.Annotated[list[Value], "FixedSize(4)"]: + ... + def two_to_two_particle_scattering_inverse(self, p1: Value, p2: Value, pa: Value, pb: Value) -> typing.Annotated[list[Value], "FixedSize(4)"]: + ... + def uniform_invariant(self, r: Value, s_min: Value, s_max: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... + def uniform_invariant_inverse(self, s: Value, s_min: Value, s_max: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... + def unsqueeze(self, input: Value) -> Value: + ... + def unstack(self, in1: Value) -> list[Value]: + ... + def unstack_sizes(self, in1: Value) -> list[Value]: + ... + def unweight(self, weights: Value, max_weight: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... + def vegas_forward(self, input: Value, grid: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... + def vegas_histogram(self, input: Value, weights: Value, bin_count: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... + def vegas_inverse(self, input: Value, grid: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... + def x1x2_to_r(self, x1: Value, x2: Value, s_lab: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... +class FunctionGenerator: + def __call__(self, *args, **kwargs): + ... + def __init__(self, name: str, arg_types: NamedTypes, return_types: NamedTypes) -> None: + ... + @typing.overload + def build_function(self, builder: FunctionBuilder, args: collections.abc.Sequence[Value]) -> NamedValues: + ... + @typing.overload + def build_function(self, builder: FunctionBuilder, args: NamedValues) -> NamedValues: + ... + def function(self) -> Function: + ... +class FunctionRuntime: + def __call__(self, *args): + ... + @typing.overload + def __init__(self, function: Function) -> None: + ... + @typing.overload + def __init__(self, function: Function, context: Context) -> None: + ... + def call(self, arg0: collections.abc.Sequence[typing.Any]) -> list[Tensor]: + ... + def call_backward(self, arg0: collections.abc.Sequence[typing.Any], arg1: collections.abc.Sequence[typing.Any], arg2: collections.abc.Sequence[bool]) -> tuple[list[madspace._madspace_py.Tensor | None], list[tuple[str, madspace._madspace_py.Tensor | None]]]: + ... + def call_with_grad(self, arg0: collections.abc.Sequence[typing.Any], arg1: collections.abc.Sequence[bool]) -> tuple[list[Tensor], list[madspace._madspace_py.Tensor | None], list[bool]]: + ... +class GeneratorConfig: + verbosity: Verbosity + write_live_data: bool + def __init__(self) -> None: + ... + @property + def combine_thread_count(self) -> int: + ... + @combine_thread_count.setter + def combine_thread_count(self, arg0: typing.SupportsInt) -> None: + ... + @property + def cpu_batch_size(self) -> int: + ... + @cpu_batch_size.setter + def cpu_batch_size(self, arg0: typing.SupportsInt) -> None: + ... + @property + def freeze_max_weight_after(self) -> int: + ... @freeze_max_weight_after.setter - def freeze_max_weight_after(self, arg0: typing.SupportsInt) -> None: ... + def freeze_max_weight_after(self, arg0: typing.SupportsInt) -> None: + ... + @property + def gpu_batch_size(self) -> int: + ... + @gpu_batch_size.setter + def gpu_batch_size(self, arg0: typing.SupportsInt) -> None: + ... @property - def max_batch_size(self) -> int: ... + def max_batch_size(self) -> int: + ... @max_batch_size.setter - def max_batch_size(self, arg0: typing.SupportsInt) -> None: ... + def max_batch_size(self, arg0: typing.SupportsInt) -> None: + ... @property - def max_overweight_truncation(self) -> float: ... + def max_overweight_truncation(self) -> float: + ... @max_overweight_truncation.setter - def max_overweight_truncation(self, arg0: typing.SupportsFloat) -> None: ... + def max_overweight_truncation(self, arg0: typing.SupportsFloat) -> None: + ... @property - def optimization_patience(self) -> int: ... + def optimization_patience(self) -> int: + ... @optimization_patience.setter - def optimization_patience(self, arg0: typing.SupportsInt) -> None: ... + def optimization_patience(self, arg0: typing.SupportsInt) -> None: + ... @property - def optimization_threshold(self) -> float: ... + def optimization_threshold(self) -> float: + ... @optimization_threshold.setter - def optimization_threshold(self, arg0: typing.SupportsFloat) -> None: ... + def optimization_threshold(self, arg0: typing.SupportsFloat) -> None: + ... @property - def start_batch_size(self) -> int: ... + def start_batch_size(self) -> int: + ... @start_batch_size.setter - def start_batch_size(self, arg0: typing.SupportsInt) -> None: ... + def start_batch_size(self, arg0: typing.SupportsInt) -> None: + ... @property - def survey_max_iters(self) -> int: ... + def survey_max_iters(self) -> int: + ... @survey_max_iters.setter - def survey_max_iters(self, arg0: typing.SupportsInt) -> None: ... + def survey_max_iters(self, arg0: typing.SupportsInt) -> None: + ... @property - def survey_min_iters(self) -> int: ... + def survey_min_iters(self) -> int: + ... @survey_min_iters.setter - def survey_min_iters(self, arg0: typing.SupportsInt) -> None: ... + def survey_min_iters(self, arg0: typing.SupportsInt) -> None: + ... @property - def survey_target_precision(self) -> float: ... + def survey_target_precision(self) -> float: + ... @survey_target_precision.setter - def survey_target_precision(self, arg0: typing.SupportsFloat) -> None: ... + def survey_target_precision(self, arg0: typing.SupportsFloat) -> None: + ... @property - def target_count(self) -> int: ... + def target_count(self) -> int: + ... @target_count.setter - def target_count(self, arg0: typing.SupportsInt) -> None: ... + def target_count(self, arg0: typing.SupportsInt) -> None: + ... @property - def vegas_damping(self) -> float: ... + def vegas_damping(self) -> float: + ... @vegas_damping.setter - def vegas_damping(self, arg0: typing.SupportsFloat) -> None: ... - -class EventGeneratorHistogram: - @property - def bin_errors(self) -> list[float]: ... - @property - def bin_values(self) -> list[float]: ... - @property - def max(self) -> float: ... - @property - def min(self) -> float: ... - @property - def name(self) -> str: ... - -class EventGeneratorStatus: + def vegas_damping(self, arg0: typing.SupportsFloat) -> None: + ... +class GeneratorStatus: done: bool - def __init__(self) -> None: ... + name: str + optimized: bool + def __init__(self) -> None: + ... @property - def count(self) -> int: ... + def count(self) -> int: + ... @count.setter - def count(self, arg0: typing.SupportsInt) -> None: ... + def count(self, arg0: typing.SupportsInt) -> None: + ... @property - def count_after_cuts(self) -> int: ... + def count_after_cuts(self) -> int: + ... @count_after_cuts.setter - def count_after_cuts(self, arg0: typing.SupportsInt) -> None: ... + def count_after_cuts(self, arg0: typing.SupportsInt) -> None: + ... @property - def count_after_cuts_opt(self) -> int: ... + def count_after_cuts_opt(self) -> int: + ... @count_after_cuts_opt.setter - def count_after_cuts_opt(self, arg0: typing.SupportsInt) -> None: ... + def count_after_cuts_opt(self, arg0: typing.SupportsInt) -> None: + ... @property - def count_opt(self) -> int: ... + def count_opt(self) -> int: + ... @count_opt.setter - def count_opt(self, arg0: typing.SupportsInt) -> None: ... + def count_opt(self, arg0: typing.SupportsInt) -> None: + ... @property - def count_target(self) -> float: ... + def count_target(self) -> float: + ... @count_target.setter - def count_target(self, arg0: typing.SupportsFloat) -> None: ... + def count_target(self, arg0: typing.SupportsFloat) -> None: + ... @property - def count_unweighted(self) -> float: ... + def count_unweighted(self) -> float: + ... @count_unweighted.setter - def count_unweighted(self, arg0: typing.SupportsFloat) -> None: ... + def count_unweighted(self, arg0: typing.SupportsFloat) -> None: + ... @property - def error(self) -> float: ... + def error(self) -> float: + ... @error.setter - def error(self, arg0: typing.SupportsFloat) -> None: ... + def error(self, arg0: typing.SupportsFloat) -> None: + ... @property - def index(self) -> int: ... - @index.setter - def index(self, arg0: typing.SupportsInt) -> None: ... - @property - def iterations(self) -> int: ... + def iterations(self) -> int: + ... @iterations.setter - def iterations(self, arg0: typing.SupportsInt) -> None: ... + def iterations(self, arg0: typing.SupportsInt) -> None: + ... @property - def mean(self) -> float: ... + def mean(self) -> float: + ... @mean.setter - def mean(self, arg0: typing.SupportsFloat) -> None: ... + def mean(self, arg0: typing.SupportsFloat) -> None: + ... @property - def rel_std_dev(self) -> float: ... + def rel_std_dev(self) -> float: + ... @rel_std_dev.setter - def rel_std_dev(self, arg0: typing.SupportsFloat) -> None: ... - -class EventGeneratorVerbosity: - """ - Members: - - silent - - log - - pretty - """ - - __members__: typing.ClassVar[ - dict[str, EventGeneratorVerbosity] - ] # value = {'silent': , 'log': , 'pretty': } - log: typing.ClassVar[ - EventGeneratorVerbosity - ] # value = - pretty: typing.ClassVar[ - EventGeneratorVerbosity - ] # value = - silent: typing.ClassVar[ - EventGeneratorVerbosity - ] # value = - def __eq__(self, other: typing.Any) -> bool: ... - def __getstate__(self) -> int: ... - def __hash__(self) -> int: ... - def __index__(self) -> int: ... - @typing.overload - def __init__(self, value: typing.SupportsInt) -> None: ... - @typing.overload - def __init__(self, name: str) -> None: ... - def __int__(self) -> int: ... - def __ne__(self, other: typing.Any) -> bool: ... - def __repr__(self) -> str: ... - def __setstate__(self, state: typing.SupportsInt) -> None: ... - def __str__(self) -> str: ... - @property - def name(self) -> str: ... - @property - def value(self) -> int: ... - -class FastRamboMapping(Mapping): - def __init__(self, n_particles: typing.SupportsInt, massless: bool) -> None: ... - -class Flow(Mapping): - def __init__( - self, - input_dim: typing.SupportsInt, - condition_dim: typing.SupportsInt = 0, - prefix: str = "", - bin_count: typing.SupportsInt = 10, - subnet_hidden_dim: typing.SupportsInt = 32, - subnet_layers: typing.SupportsInt = 3, - subnet_activation: MLP.Activation = MLP.Activation.Activation.leaky_relu, - invert_spline: bool = True, - ) -> None: ... - def condition_dim(self) -> int: ... - def initialize_from_vegas(self, context: Context, grid_name: str) -> None: ... - def initialize_globals(self, context: Context) -> None: ... - def input_dim(self) -> int: ... - -class Function: - @staticmethod - def load(file: str) -> Function: ... - def __call__(self, *args): ... - def __repr__(self) -> str: ... - def __str__(self) -> str: ... - def save(self, file: str) -> None: ... + def rel_std_dev(self, arg0: typing.SupportsFloat) -> None: + ... + @property + def subprocess(self) -> int: + ... + @subprocess.setter + def subprocess(self, arg0: typing.SupportsInt) -> None: + ... +class HistItem: + def __init__(self, observable: Observable, min: typing.SupportsFloat, max: typing.SupportsFloat, bin_count: typing.SupportsInt) -> None: + ... @property - def globals(self) -> dict[str, Value]: ... + def bin_count(self) -> int: + ... @property - def inputs(self) -> list[Value]: ... + def max(self) -> float: + ... @property - def instructions(self) -> list[InstructionCall]: ... + def min(self) -> float: + ... @property - def locals(self) -> list[Value]: ... + def observable(self) -> Observable: + ... +class Histogram: @property - def outputs(self) -> list[Value]: ... - -class FunctionBuilder: - def __init__( - self, - input_types: collections.abc.Sequence[Type], - output_types: collections.abc.Sequence[Type], - ) -> None: ... - def add(self, in1: Value, in2: Value) -> Value: ... - def add_int(self, in1: Value, in2: Value) -> Value: ... - def apply_subchannel_weights( - self, - channel_weights_in: Value, - subchannel_weights: Value, - channel_indices: Value, - subchannel_indices: Value, - ) -> Value: ... - def argsort(self, input: Value) -> Value: ... - def batch_cat( - self, args: collections.abc.Sequence[Value] - ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... - def batch_gather(self, indices: Value, values: Value) -> Value: ... - def batch_scatter(self, indices: Value, target: Value, source: Value) -> Value: ... - def batch_size(self, args: collections.abc.Sequence[Value]) -> Value: ... - def batch_split(self, in1: Value, counts: Value) -> list[Value]: ... - def boost_beam(self, p1: Value, x1: Value, x2: Value) -> Value: ... - def boost_beam_inverse(self, p1: Value, x1: Value, x2: Value) -> Value: ... - def breit_wigner_invariant( - self, r: Value, mass: Value, width: Value, s_min: Value, s_max: Value - ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... - def breit_wigner_invariant_inverse( - self, s: Value, mass: Value, width: Value, s_min: Value, s_max: Value - ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... - def cat(self, args: collections.abc.Sequence[Value]) -> Value: ... - def chili_forward( - self, r: Value, e_cm: Value, m_out: Value, pt_min: Value, y_max: Value - ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... - def chili_inverse( - self, p_ext: Value, e_cm: Value, m_out: Value, pt_min: Value, y_max: Value - ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... - def collect_channel_weights( - self, amp2: Value, channel_indices: Value, channel_count: Value - ) -> Value: ... - def com_p_in( - self, e_cm: Value - ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... - def cut_all(self, obs: Value, min: Value, max: Value) -> Value: ... - def cut_any(self, obs: Value, min: Value, max: Value) -> Value: ... - def cut_one(self, obs: Value, min: Value, max: Value) -> Value: ... - def cut_unphysical(self, w_in: Value, p: Value, x1: Value, x2: Value) -> Value: ... - def diff_cross_section( - self, - x1: Value, - x2: Value, - pdf1: Value, - pdf2: Value, - matrix_element: Value, - e_cm2: Value, - ) -> Value: ... - def discrete_histogram( - self, input: Value, weights: Value, option_count: Value - ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... - def elu(self, in1: Value) -> Value: ... - def fast_rambo_massive( - self, r: Value, e_cm: Value, masses: Value, p0: Value - ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... - def fast_rambo_massive_com( - self, r: Value, e_cm: Value, masses: Value - ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... - def fast_rambo_massive_inverse( - self, p_out: Value, e_cm: Value, masses: Value - ) -> typing.Annotated[list[Value], "FixedSize(3)"]: ... - def fast_rambo_massless( - self, r: Value, e_cm: Value, p0: Value - ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... - def fast_rambo_massless_com( - self, r: Value, e_cm: Value - ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... - def fast_rambo_massless_inverse( - self, p_out: Value, e_cm: Value - ) -> typing.Annotated[list[Value], "FixedSize(3)"]: ... - def full(self, args: collections.abc.Sequence[Value]) -> Value: ... - def function(self) -> Function: ... - def gather(self, index: Value, choices: Value) -> Value: ... - def gather_int(self, index: Value, choices: Value) -> Value: ... - def gelu(self, in1: Value) -> Value: ... - def get_global( - self, - name: str, - dtype: DataType, - shape: collections.abc.Sequence[typing.SupportsInt], - ) -> Value: ... - def histogram( - self, input: Value, weights: Value, min: Value, max: Value, bin_count: Value - ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... - def input(self, index: typing.SupportsInt) -> Value: ... - def input_range( - self, start_index: typing.SupportsInt, end_index: typing.SupportsInt - ) -> list[Value]: ... - def interpolate_alpha_s( - self, q2: Value, grid_logq2: Value, grid_coeffs: Value - ) -> Value: ... - def interpolate_pdf( - self, - x: Value, - q2: Value, - pid_indices: Value, - grid_logx: Value, - grid_logq2: Value, - grid_coeffs: Value, - ) -> Value: ... - def invariants_from_momenta(self, p_ext: Value, factors: Value) -> Value: ... - def leaky_relu(self, in1: Value) -> Value: ... - def matmul(self, x: Value, weight: Value, bias: Value) -> Value: ... - def matrix_element(self, args: collections.abc.Sequence[Value]) -> list[Value]: ... - def max(self, in1: Value, in2: Value) -> Value: ... - def min(self, in1: Value, in2: Value) -> Value: ... - def mirror_momenta(self, p_ext: Value, mirror: Value) -> Value: ... - def momenta_to_x1x2( - self, p_ext: Value, e_cm: Value - ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... - def mul(self, in1: Value, in2: Value) -> Value: ... - def nonzero(self, input: Value) -> Value: ... - def obs_delta_eta(self, p1: Value, p2: Value) -> Value: ... - def obs_delta_phi(self, p1: Value, p2: Value) -> Value: ... - def obs_delta_r(self, p1: Value, p2: Value) -> Value: ... - def obs_e(self, p: Value) -> Value: ... - def obs_eta(self, p: Value) -> Value: ... - def obs_eta_abs(self, p: Value) -> Value: ... - def obs_mass(self, p: Value) -> Value: ... - def obs_p_mag(self, p: Value) -> Value: ... - def obs_phi(self, p: Value) -> Value: ... - def obs_pt(self, p: Value) -> Value: ... - def obs_px(self, p: Value) -> Value: ... - def obs_py(self, p: Value) -> Value: ... - def obs_pz(self, p: Value) -> Value: ... - def obs_sqrt_s(self, p_ext: Value) -> Value: ... - def obs_theta(self, p: Value) -> Value: ... - def obs_y(self, p: Value) -> Value: ... - def obs_y_abs(self, p: Value) -> Value: ... - def offset_indices( - self, batch_sizes_offset: Value, batch_sizes_out: Value - ) -> Value: ... - def one_hot(self, index: Value, option_count: Value) -> Value: ... - def output(self, index: typing.SupportsInt, value: Value) -> None: ... - def output_range( - self, start_index: typing.SupportsInt, values: collections.abc.Sequence[Value] - ) -> None: ... - def permute_momenta( - self, momenta: Value, permutations: Value, index: Value - ) -> Value: ... - def pop(self, in1: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... - def product(self, values: collections.abc.Sequence[Value]) -> Value: ... - def pt_eta_phi_x(self, p_ext: Value, x1: Value, x2: Value) -> Value: ... - def r_to_x1x2( - self, r: Value, s_hat: Value, s_lab: Value - ) -> typing.Annotated[list[Value], "FixedSize(3)"]: ... - def random(self, batch_size: Value, count: Value) -> Value: ... - def reduce_product(self, in1: Value) -> Value: ... - def reduce_sum(self, in1: Value) -> Value: ... - def reduce_sum_vector(self, in1: Value) -> Value: ... - def relu(self, in1: Value) -> Value: ... - def rqs_find_bin( - self, input: Value, in_sizes: Value, out_sizes: Value, derivatives: Value - ) -> Value: ... - def rqs_forward( - self, input: Value, condition: Value - ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... - def rqs_inverse( - self, input: Value, condition: Value - ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... - def rqs_reshape( - self, input: Value, bin_count: Value - ) -> typing.Annotated[list[Value], "FixedSize(3)"]: ... - def s23_min_max( - self, pa: Value, pb: Value, p3: Value, t1_abs: Value, m1: Value, m2: Value - ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... - def s23_value_and_min_max( - self, pa: Value, pb: Value, p3: Value, t1_abs: Value, p1: Value, p2: Value - ) -> typing.Annotated[list[Value], "FixedSize(3)"]: ... - def sample_discrete( - self, r: Value, option_count: Value - ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... - def sample_discrete_inverse( - self, index: Value, option_count: Value - ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... - def sample_discrete_probs( - self, r: Value, probs: Value - ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... - def sample_discrete_probs_inverse( - self, index: Value, probs: Value - ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... - def scale_half_transverse_mass(self, momenta: Value) -> Value: ... - def scale_partonic_energy(self, momenta: Value) -> Value: ... - def scale_transverse_energy(self, momenta: Value) -> Value: ... - def scale_transverse_mass(self, momenta: Value) -> Value: ... - def sde2_channel_weights( - self, invariants: Value, masses: Value, widths: Value, indices: Value - ) -> Value: ... - def select(self, input: Value, indices: Value) -> Value: ... - def select_int(self, input: Value, indices: Value) -> Value: ... - def select_vector(self, input: Value, indices: Value) -> Value: ... - def sigmoid(self, in1: Value) -> Value: ... - def softmax(self, input: Value) -> Value: ... - def softmax_prior(self, input: Value, prior: Value) -> Value: ... - def softplus(self, in1: Value) -> Value: ... - def sqrt(self, in1: Value) -> Value: ... - def square(self, in1: Value) -> Value: ... - def squeeze(self, input: Value) -> Value: ... - def stable_invariant( - self, r: Value, mass: Value, s_min: Value, s_max: Value - ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... - def stable_invariant_inverse( - self, s: Value, mass: Value, s_min: Value, s_max: Value - ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... - def stable_invariant_nu( - self, r: Value, mass: Value, nu: Value, s_min: Value, s_max: Value - ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... - def stable_invariant_nu_inverse( - self, s: Value, mass: Value, nu: Value, s_min: Value, s_max: Value - ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... - def stack(self, args: collections.abc.Sequence[Value]) -> Value: ... - def sub(self, in1: Value, in2: Value) -> Value: ... - def subchannel_weights( - self, - invariants: Value, - masses: Value, - widths: Value, - indices: Value, - on_shell: Value, - group_sizes: Value, - ) -> Value: ... - def t_inv_min_max( - self, pa: Value, pb: Value, m1: Value, m2: Value - ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... - def t_inv_value_and_min_max( - self, pa: Value, pb: Value, p1: Value, p2: Value - ) -> typing.Annotated[list[Value], "FixedSize(3)"]: ... - def three_body_decay( - self, - r_e1: Value, - r_e2: Value, - r_phi: Value, - r_cos_theta: Value, - r_beta: Value, - m0: Value, - m1: Value, - m2: Value, - m3: Value, - p0: Value, - ) -> typing.Annotated[list[Value], "FixedSize(4)"]: ... - def three_body_decay_com( - self, - r_e1: Value, - r_e2: Value, - r_phi: Value, - r_cos_theta: Value, - r_beta: Value, - m0: Value, - m1: Value, - m2: Value, - m3: Value, - ) -> typing.Annotated[list[Value], "FixedSize(4)"]: ... - def three_body_decay_com_inverse( - self, p1: Value, p2: Value, p3: Value - ) -> typing.Annotated[list[Value], "FixedSize(10)"]: ... - def three_body_decay_inverse( - self, p1: Value, p2: Value, p3: Value - ) -> typing.Annotated[list[Value], "FixedSize(11)"]: ... - def two_body_decay( - self, - r_phi: Value, - r_cos_theta: Value, - m0: Value, - m1: Value, - m2: Value, - p0: Value, - ) -> typing.Annotated[list[Value], "FixedSize(3)"]: ... - def two_body_decay_com( - self, r_phi: Value, r_cos_theta: Value, m0: Value, m1: Value, m2: Value - ) -> typing.Annotated[list[Value], "FixedSize(3)"]: ... - def two_body_decay_com_inverse( - self, p1: Value, p2: Value - ) -> typing.Annotated[list[Value], "FixedSize(6)"]: ... - def two_body_decay_inverse( - self, p1: Value, p2: Value - ) -> typing.Annotated[list[Value], "FixedSize(7)"]: ... - def two_to_three_particle_scattering( - self, - phi_choice: Value, - pa: Value, - pb: Value, - p3: Value, - s23: Value, - t1_abs: Value, - m1: Value, - m2: Value, - ) -> typing.Annotated[list[Value], "FixedSize(3)"]: ... - def two_to_three_particle_scattering_inverse( - self, - p1: Value, - p2: Value, - p3: Value, - pa: Value, - pb: Value, - t1_abs: Value, - s23: Value, - ) -> typing.Annotated[list[Value], "FixedSize(4)"]: ... - def two_to_two_particle_scattering( - self, r_phi: Value, pa: Value, pb: Value, t_abs: Value, m1: Value, m2: Value - ) -> typing.Annotated[list[Value], "FixedSize(3)"]: ... - def two_to_two_particle_scattering_com( - self, r_phi: Value, pa: Value, pb: Value, t_abs: Value, m1: Value, m2: Value - ) -> typing.Annotated[list[Value], "FixedSize(3)"]: ... - def two_to_two_particle_scattering_com_inverse( - self, p1: Value, p2: Value, pa: Value, pb: Value - ) -> typing.Annotated[list[Value], "FixedSize(4)"]: ... - def two_to_two_particle_scattering_inverse( - self, p1: Value, p2: Value, pa: Value, pb: Value - ) -> typing.Annotated[list[Value], "FixedSize(4)"]: ... - def uniform_invariant( - self, r: Value, s_min: Value, s_max: Value - ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... - def uniform_invariant_inverse( - self, s: Value, s_min: Value, s_max: Value - ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... - def unsqueeze(self, input: Value) -> Value: ... - def unstack(self, in1: Value) -> list[Value]: ... - def unstack_sizes(self, in1: Value) -> list[Value]: ... - def unweight( - self, weights: Value, max_weight: Value - ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... - def vegas_forward( - self, input: Value, grid: Value - ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... - def vegas_histogram( - self, input: Value, weights: Value, bin_count: Value - ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... - def vegas_inverse( - self, input: Value, grid: Value - ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... - def x1x2_to_r( - self, x1: Value, x2: Value, s_lab: Value - ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... - -class FunctionGenerator: - def __call__(self, *args): ... - def __init__( - self, - name: str, - arg_types: collections.abc.Sequence[Type], - return_types: collections.abc.Sequence[Type], - ) -> None: ... - def build_function( - self, builder: FunctionBuilder, args: collections.abc.Sequence[Value] - ) -> list[Value]: ... - def function(self) -> Function: ... - -class FunctionRuntime: - def __call__(self, *args): ... - @typing.overload - def __init__(self, function: Function) -> None: ... - @typing.overload - def __init__(self, function: Function, context: Context) -> None: ... - def call(self, arg0: collections.abc.Sequence[typing.Any]) -> list[Tensor]: ... - def call_backward( - self, - arg0: collections.abc.Sequence[typing.Any], - arg1: collections.abc.Sequence[typing.Any], - arg2: collections.abc.Sequence[bool], - ) -> tuple[ - list[madspace._madspace_py.Tensor | None], - list[tuple[str, madspace._madspace_py.Tensor | None]], - ]: ... - def call_with_grad( - self, - arg0: collections.abc.Sequence[typing.Any], - arg1: collections.abc.Sequence[bool], - ) -> tuple[list[Tensor], list[madspace._madspace_py.Tensor | None], list[bool]]: ... - -class HistItem: - def __init__( - self, - observable: Observable, - min: typing.SupportsFloat, - max: typing.SupportsFloat, - bin_count: typing.SupportsInt, - ) -> None: ... + def bin_errors(self) -> list[float]: + ... @property - def bin_count(self) -> int: ... + def bin_values(self) -> list[float]: + ... @property - def max(self) -> float: ... + def max(self) -> float: + ... @property - def min(self) -> float: ... + def min(self) -> float: + ... @property - def observable(self) -> Observable: ... - + def name(self) -> str: + ... class Instruction: - def __str__(self) -> str: ... + def __str__(self) -> str: + ... @property - def name(self) -> str: ... + def name(self) -> str: + ... @property - def opcode(self) -> int: ... - + def opcode(self) -> int: + ... class InstructionCall: - def __repr__(self) -> str: ... - def __str__(self) -> str: ... + def __repr__(self) -> str: + ... + def __str__(self) -> str: + ... @property - def inputs(self) -> list[Value]: ... + def inputs(self) -> list[Value]: + ... @property - def instruction(self) -> Instruction: ... + def instruction(self) -> Instruction: + ... @property - def outputs(self) -> list[Value]: ... - + def outputs(self) -> list[Value]: + ... class Integrand(FunctionGenerator): - matrix_element_inputs: typing.ClassVar[ - list - ] # value = [, , , , , ] - matrix_element_outputs: typing.ClassVar[ - list - ] # value = [, , , , ] + drop_cuts_and_rescale: typing.ClassVar[int] = 8192 + exclude_adaptive_and_chan_weight: typing.ClassVar[int] = 4096 + matrix_element_inputs: typing.ClassVar[list] # value = [, , , , , ] + matrix_element_outputs: typing.ClassVar[list] # value = [, , , , ] return_chan_weights: typing.ClassVar[int] = 256 return_channel: typing.ClassVar[int] = 128 return_cwnet_input: typing.ClassVar[int] = 512 @@ -1100,425 +1069,419 @@ class Integrand(FunctionGenerator): return_x1_x2: typing.ClassVar[int] = 8 sample: typing.ClassVar[int] = 1 unweight: typing.ClassVar[int] = 2 - def __init__( - self, - mapping: PhaseSpaceMapping, - diff_xs: DifferentialCrossSection, - adaptive_map: ( - None | madspace._madspace_py.VegasMapping | madspace._madspace_py.Flow - ) = None, - discrete_before: ( - None - | madspace._madspace_py.DiscreteSampler - | madspace._madspace_py.DiscreteFlow - ) = None, - discrete_after: ( - None - | madspace._madspace_py.DiscreteSampler - | madspace._madspace_py.DiscreteFlow - ) = None, - pdf_grid: madspace._madspace_py.PdfGrid | None = None, - energy_scale: madspace._madspace_py.EnergyScale | None = None, - prop_chan_weights: madspace._madspace_py.PropagatorChannelWeights | None = None, - subchan_weights: madspace._madspace_py.SubchannelWeights | None = None, - chan_weight_net: madspace._madspace_py.ChannelWeightNetwork | None = None, - chan_weight_remap: collections.abc.Sequence[typing.SupportsInt] = [], - remapped_chan_count: typing.SupportsInt = 0, - flags: typing.SupportsInt = 0, - channel_indices: collections.abc.Sequence[typing.SupportsInt] = [], - active_flavors: collections.abc.Sequence[typing.SupportsInt] = [], - ) -> None: ... - def adaptive_map( - self, - ) -> None | madspace._madspace_py.VegasMapping | madspace._madspace_py.Flow: ... - def chan_weight_net(self) -> madspace._madspace_py.ChannelWeightNetwork | None: ... - def diff_xs(self) -> DifferentialCrossSection: ... - def discrete_after( - self, - ) -> ( - None - | madspace._madspace_py.DiscreteSampler - | madspace._madspace_py.DiscreteFlow - ): ... - def discrete_before( - self, - ) -> ( - None - | madspace._madspace_py.DiscreteSampler - | madspace._madspace_py.DiscreteFlow - ): ... - def energy_scale(self) -> madspace._madspace_py.EnergyScale | None: ... - def flags(self) -> int: ... - def latent_dims(self) -> tuple[list[int], list[bool]]: ... - def mapping(self) -> PhaseSpaceMapping: ... - def particle_count(self) -> int: ... - def prop_chan_weights( - self, - ) -> madspace._madspace_py.PropagatorChannelWeights | None: ... - def random_dim(self) -> int: ... - def vegas_grid_name(self) -> str | None: ... - + def __init__(self, mapping: PhaseSpaceMapping, diff_xs: DifferentialCrossSection, adaptive_map: None | madspace._madspace_py.VegasMapping | madspace._madspace_py.Flow = None, discrete_before: None | madspace._madspace_py.DiscreteSampler | madspace._madspace_py.DiscreteFlow = None, discrete_after: None | madspace._madspace_py.DiscreteSampler | madspace._madspace_py.DiscreteFlow = None, pdf_grid: madspace._madspace_py.PdfGrid | None = None, energy_scale: madspace._madspace_py.EnergyScale | None = None, prop_chan_weights: madspace._madspace_py.PropagatorChannelWeights | None = None, subchan_weights: madspace._madspace_py.SubchannelWeights | None = None, chan_weight_net: madspace._madspace_py.ChannelWeightNetwork | None = None, chan_weight_remap: collections.abc.Sequence[typing.SupportsInt] = [], remapped_chan_count: typing.SupportsInt = 0, flags: typing.SupportsInt = 0, channel_indices: collections.abc.Sequence[typing.SupportsInt] = [], active_flavors: collections.abc.Sequence[typing.SupportsInt] = [], flavor_remap: collections.abc.Sequence[typing.SupportsInt] = [], flavor_factors: collections.abc.Sequence[typing.SupportsFloat] = []) -> None: + ... + def adaptive_map(self) -> None | madspace._madspace_py.VegasMapping | madspace._madspace_py.Flow: + ... + def chan_weight_net(self) -> madspace._madspace_py.ChannelWeightNetwork | None: + ... + def diff_xs(self) -> DifferentialCrossSection: + ... + def discrete_after(self) -> None | madspace._madspace_py.DiscreteSampler | madspace._madspace_py.DiscreteFlow: + ... + def discrete_before(self) -> None | madspace._madspace_py.DiscreteSampler | madspace._madspace_py.DiscreteFlow: + ... + def energy_scale(self) -> madspace._madspace_py.EnergyScale | None: + ... + def flags(self) -> int: + ... + def latent_dims(self) -> tuple[list[int], list[bool]]: + ... + def mapping(self) -> PhaseSpaceMapping: + ... + def particle_count(self) -> int: + ... + def prop_chan_weights(self) -> madspace._madspace_py.PropagatorChannelWeights | None: + ... + def random_dim(self) -> int: + ... + def vegas_grid_name(self) -> str | None: + ... class IntegrandProbability(FunctionGenerator): - def __init__(self, integrand: Integrand) -> None: ... - + def __init__(self, integrand: Integrand) -> None: + ... class Invariant(Mapping): - def __init__( - self, - power: typing.SupportsFloat = 0.0, - mass: typing.SupportsFloat = 0.0, - width: typing.SupportsFloat = 0.0, - ) -> None: ... - + def __init__(self, power: typing.SupportsFloat = 0.0, mass: typing.SupportsFloat = 0.0, width: typing.SupportsFloat = 0.0) -> None: + ... class LHECompleter: - def __init__( - self, - subproc_args: collections.abc.Sequence[SubprocArgs], - bw_cutoff: typing.SupportsFloat, - ) -> None: ... - def complete_event_data( - self, - event: LHEEvent, - subprocess_index: typing.SupportsInt, - diagram_index: typing.SupportsInt, - color_index: typing.SupportsInt, - flavor_index: typing.SupportsInt, - helicity_index: typing.SupportsInt, - ) -> None: ... - @property - def max_particle_count(self) -> int: ... - + @staticmethod + def load(file: str) -> LHECompleter: + ... + def __init__(self, subproc_args: collections.abc.Sequence[SubprocArgs], bw_cutoff: typing.SupportsFloat) -> None: + ... + def save(self, file: str) -> None: + ... + @property + def max_particle_count(self) -> int: + ... class LHEEvent: - def __init__( - self, - process_id: typing.SupportsInt = 0, - weight: typing.SupportsFloat = 0.0, - scale: typing.SupportsFloat = 0.0, - alpha_qed: typing.SupportsFloat = 0.0, - alpha_qcd: typing.SupportsFloat = 0.0, - particles: collections.abc.Sequence[LHEParticle] = [], - ) -> None: ... - @property - def alpha_qcd(self) -> int: ... + def __init__(self, process_id: typing.SupportsInt = 0, weight: typing.SupportsFloat = 0.0, scale: typing.SupportsFloat = 0.0, alpha_qed: typing.SupportsFloat = 0.0, alpha_qcd: typing.SupportsFloat = 0.0, particles: collections.abc.Sequence[LHEParticle] = []) -> None: + ... + @property + def alpha_qcd(self) -> int: + ... @alpha_qcd.setter - def alpha_qcd(self, arg0: typing.SupportsInt) -> None: ... + def alpha_qcd(self, arg0: typing.SupportsInt) -> None: + ... @property - def alpha_qed(self) -> float: ... + def alpha_qed(self) -> float: + ... @alpha_qed.setter - def alpha_qed(self, arg0: typing.SupportsFloat) -> None: ... + def alpha_qed(self, arg0: typing.SupportsFloat) -> None: + ... @property - def particles(self) -> list[LHEParticle]: ... + def particles(self) -> list[LHEParticle]: + ... @particles.setter - def particles(self, arg0: collections.abc.Sequence[LHEParticle]) -> None: ... + def particles(self, arg0: collections.abc.Sequence[LHEParticle]) -> None: + ... @property - def process_id(self) -> int: ... + def process_id(self) -> int: + ... @process_id.setter - def process_id(self, arg0: typing.SupportsInt) -> None: ... + def process_id(self, arg0: typing.SupportsInt) -> None: + ... @property - def scale(self) -> float: ... + def scale(self) -> float: + ... @scale.setter - def scale(self, arg0: typing.SupportsFloat) -> None: ... + def scale(self, arg0: typing.SupportsFloat) -> None: + ... @property - def weight(self) -> float: ... + def weight(self) -> float: + ... @weight.setter - def weight(self, arg0: typing.SupportsFloat) -> None: ... - + def weight(self, arg0: typing.SupportsFloat) -> None: + ... class LHEFileWriter: - def __init__(self, file_name: str, meta: LHEMeta) -> None: ... - def write(self, event: LHEEvent) -> None: ... - def write_string(self, str: str) -> None: ... - + def __init__(self, file_name: str, meta: LHEMeta) -> None: + ... + def write(self, event: LHEEvent) -> None: + ... + def write_string(self, str: str) -> None: + ... class LHEHeader: content: str escape_content: bool name: str - def __init__( - self, name: str = "", content: str = "", escape_content: bool = False - ) -> None: ... - + def __init__(self, name: str = '', content: str = '', escape_content: bool = False) -> None: + ... class LHEMeta: - def __init__( - self, - beam1_pdg_id: typing.SupportsInt = 0, - beam2_pdg_id: typing.SupportsInt = 0, - beam1_energy: typing.SupportsFloat = 0.0, - beam2_energy: typing.SupportsFloat = 0.0, - beam1_pdf_authors: typing.SupportsInt = 0, - beam2_pdf_authors: typing.SupportsInt = 0, - beam1_pdf_id: typing.SupportsInt = 0, - beam2_pdf_id: typing.SupportsInt = 0, - weight_mode: typing.SupportsInt = 0, - processes: collections.abc.Sequence[LHEProcess] = [], - headers: collections.abc.Sequence[LHEHeader] = [], - ) -> None: ... - @property - def beam1_energy(self) -> float: ... + def __init__(self, beam1_pdg_id: typing.SupportsInt = 0, beam2_pdg_id: typing.SupportsInt = 0, beam1_energy: typing.SupportsFloat = 0.0, beam2_energy: typing.SupportsFloat = 0.0, beam1_pdf_authors: typing.SupportsInt = 0, beam2_pdf_authors: typing.SupportsInt = 0, beam1_pdf_id: typing.SupportsInt = 0, beam2_pdf_id: typing.SupportsInt = 0, weight_mode: typing.SupportsInt = 0, processes: collections.abc.Sequence[LHEProcess] = [], headers: collections.abc.Sequence[LHEHeader] = []) -> None: + ... + @property + def beam1_energy(self) -> float: + ... @beam1_energy.setter - def beam1_energy(self, arg0: typing.SupportsFloat) -> None: ... + def beam1_energy(self, arg0: typing.SupportsFloat) -> None: + ... @property - def beam1_pdf_authors(self) -> int: ... + def beam1_pdf_authors(self) -> int: + ... @beam1_pdf_authors.setter - def beam1_pdf_authors(self, arg0: typing.SupportsInt) -> None: ... + def beam1_pdf_authors(self, arg0: typing.SupportsInt) -> None: + ... @property - def beam1_pdf_id(self) -> int: ... + def beam1_pdf_id(self) -> int: + ... @beam1_pdf_id.setter - def beam1_pdf_id(self, arg0: typing.SupportsInt) -> None: ... + def beam1_pdf_id(self, arg0: typing.SupportsInt) -> None: + ... @property - def beam1_pdg_id(self) -> int: ... + def beam1_pdg_id(self) -> int: + ... @beam1_pdg_id.setter - def beam1_pdg_id(self, arg0: typing.SupportsInt) -> None: ... + def beam1_pdg_id(self, arg0: typing.SupportsInt) -> None: + ... @property - def beam2_energy(self) -> float: ... + def beam2_energy(self) -> float: + ... @beam2_energy.setter - def beam2_energy(self, arg0: typing.SupportsFloat) -> None: ... + def beam2_energy(self, arg0: typing.SupportsFloat) -> None: + ... @property - def beam2_pdf_authors(self) -> int: ... + def beam2_pdf_authors(self) -> int: + ... @beam2_pdf_authors.setter - def beam2_pdf_authors(self, arg0: typing.SupportsInt) -> None: ... + def beam2_pdf_authors(self, arg0: typing.SupportsInt) -> None: + ... @property - def beam2_pdf_id(self) -> int: ... + def beam2_pdf_id(self) -> int: + ... @beam2_pdf_id.setter - def beam2_pdf_id(self, arg0: typing.SupportsInt) -> None: ... + def beam2_pdf_id(self, arg0: typing.SupportsInt) -> None: + ... @property - def beam2_pdg_id(self) -> int: ... + def beam2_pdg_id(self) -> int: + ... @beam2_pdg_id.setter - def beam2_pdg_id(self, arg0: typing.SupportsInt) -> None: ... + def beam2_pdg_id(self, arg0: typing.SupportsInt) -> None: + ... @property - def headers(self) -> list[LHEHeader]: ... + def headers(self) -> list[LHEHeader]: + ... @headers.setter - def headers(self, arg0: collections.abc.Sequence[LHEHeader]) -> None: ... + def headers(self, arg0: collections.abc.Sequence[LHEHeader]) -> None: + ... @property - def processes(self) -> list[LHEProcess]: ... + def processes(self) -> list[LHEProcess]: + ... @processes.setter - def processes(self, arg0: collections.abc.Sequence[LHEProcess]) -> None: ... + def processes(self, arg0: collections.abc.Sequence[LHEProcess]) -> None: + ... @property - def weight_mode(self) -> int: ... + def weight_mode(self) -> int: + ... @weight_mode.setter - def weight_mode(self, arg0: typing.SupportsInt) -> None: ... - + def weight_mode(self, arg0: typing.SupportsInt) -> None: + ... class LHEParticle: status_incoming: typing.ClassVar[int] = -1 status_intermediate_resonance: typing.ClassVar[int] = 2 status_outgoing: typing.ClassVar[int] = 1 - def __init__( - self, - pdg_id: typing.SupportsInt = 0, - status_code: typing.SupportsInt = 0, - mother1: typing.SupportsInt = 0, - mother2: typing.SupportsInt = 0, - color: typing.SupportsInt = 0, - anti_color: typing.SupportsInt = 0, - p_x: typing.SupportsFloat = 0.0, - p_y: typing.SupportsFloat = 0.0, - p_z: typing.SupportsFloat = 0.0, - energy: typing.SupportsFloat = 0.0, - mass: typing.SupportsFloat = 0.0, - lifetime: typing.SupportsFloat = 0.0, - spin: typing.SupportsFloat = 0.0, - ) -> None: ... - @property - def anti_color(self) -> int: ... + def __init__(self, pdg_id: typing.SupportsInt = 0, status_code: typing.SupportsInt = 0, mother1: typing.SupportsInt = 0, mother2: typing.SupportsInt = 0, color: typing.SupportsInt = 0, anti_color: typing.SupportsInt = 0, p_x: typing.SupportsFloat = 0.0, p_y: typing.SupportsFloat = 0.0, p_z: typing.SupportsFloat = 0.0, energy: typing.SupportsFloat = 0.0, mass: typing.SupportsFloat = 0.0, lifetime: typing.SupportsFloat = 0.0, spin: typing.SupportsFloat = 0.0) -> None: + ... + @property + def anti_color(self) -> int: + ... @anti_color.setter - def anti_color(self, arg0: typing.SupportsInt) -> None: ... + def anti_color(self, arg0: typing.SupportsInt) -> None: + ... @property - def color(self) -> int: ... + def color(self) -> int: + ... @color.setter - def color(self, arg0: typing.SupportsInt) -> None: ... + def color(self, arg0: typing.SupportsInt) -> None: + ... @property - def energy(self) -> float: ... + def energy(self) -> float: + ... @energy.setter - def energy(self, arg0: typing.SupportsFloat) -> None: ... + def energy(self, arg0: typing.SupportsFloat) -> None: + ... @property - def lifetime(self) -> float: ... + def lifetime(self) -> float: + ... @lifetime.setter - def lifetime(self, arg0: typing.SupportsFloat) -> None: ... + def lifetime(self, arg0: typing.SupportsFloat) -> None: + ... @property - def mass(self) -> float: ... + def mass(self) -> float: + ... @mass.setter - def mass(self, arg0: typing.SupportsFloat) -> None: ... + def mass(self, arg0: typing.SupportsFloat) -> None: + ... @property - def mother1(self) -> int: ... + def mother1(self) -> int: + ... @mother1.setter - def mother1(self, arg0: typing.SupportsInt) -> None: ... + def mother1(self, arg0: typing.SupportsInt) -> None: + ... @property - def mother2(self) -> int: ... + def mother2(self) -> int: + ... @mother2.setter - def mother2(self, arg0: typing.SupportsInt) -> None: ... + def mother2(self, arg0: typing.SupportsInt) -> None: + ... @property - def pdg_id(self) -> int: ... + def pdg_id(self) -> int: + ... @pdg_id.setter - def pdg_id(self, arg0: typing.SupportsInt) -> None: ... + def pdg_id(self, arg0: typing.SupportsInt) -> None: + ... @property - def px(self) -> float: ... + def px(self) -> float: + ... @px.setter - def px(self, arg0: typing.SupportsFloat) -> None: ... + def px(self, arg0: typing.SupportsFloat) -> None: + ... @property - def py(self) -> float: ... + def py(self) -> float: + ... @py.setter - def py(self, arg0: typing.SupportsFloat) -> None: ... + def py(self, arg0: typing.SupportsFloat) -> None: + ... @property - def pz(self) -> float: ... + def pz(self) -> float: + ... @pz.setter - def pz(self, arg0: typing.SupportsFloat) -> None: ... + def pz(self, arg0: typing.SupportsFloat) -> None: + ... @property - def spin(self) -> float: ... + def spin(self) -> float: + ... @spin.setter - def spin(self, arg0: typing.SupportsFloat) -> None: ... + def spin(self, arg0: typing.SupportsFloat) -> None: + ... @property - def status_code(self) -> int: ... + def status_code(self) -> int: + ... @status_code.setter - def status_code(self, arg0: typing.SupportsInt) -> None: ... - + def status_code(self, arg0: typing.SupportsInt) -> None: + ... class LHEProcess: @staticmethod - def __init__(*args, **kwargs) -> None: ... + def __init__(*args, **kwargs) -> None: + ... @property - def cross_section(self) -> float: ... + def cross_section(self) -> float: + ... @cross_section.setter - def cross_section(self, arg0: typing.SupportsFloat) -> None: ... + def cross_section(self, arg0: typing.SupportsFloat) -> None: + ... @property - def cross_section_error(self) -> float: ... + def cross_section_error(self) -> float: + ... @cross_section_error.setter - def cross_section_error(self, arg0: typing.SupportsFloat) -> None: ... + def cross_section_error(self, arg0: typing.SupportsFloat) -> None: + ... @property - def max_weight(self) -> float: ... + def max_weight(self) -> float: + ... @max_weight.setter - def max_weight(self, arg0: typing.SupportsFloat) -> None: ... + def max_weight(self, arg0: typing.SupportsFloat) -> None: + ... @property - def process_id(self) -> int: ... + def process_id(self) -> int: + ... @process_id.setter - def process_id(self, arg0: typing.SupportsInt) -> None: ... - + def process_id(self, arg0: typing.SupportsInt) -> None: + ... class LineRef: - def __init__(self, str: str) -> None: ... - def __repr__(self) -> str: ... - + def __init__(self, str: str) -> None: + ... + def __repr__(self) -> str: + ... class Logger: class LogLevel: """ Members: - + level_debug - + level_info - + level_warning - + level_error """ - - __members__: typing.ClassVar[ - dict[str, Logger.LogLevel] - ] # value = {'level_debug': , 'level_info': , 'level_warning': , 'level_error': } - level_debug: typing.ClassVar[ - Logger.LogLevel - ] # value = - level_error: typing.ClassVar[ - Logger.LogLevel - ] # value = + __members__: typing.ClassVar[dict[str, Logger.LogLevel]] # value = {'level_debug': , 'level_info': , 'level_warning': , 'level_error': } + level_debug: typing.ClassVar[Logger.LogLevel] # value = + level_error: typing.ClassVar[Logger.LogLevel] # value = level_info: typing.ClassVar[Logger.LogLevel] # value = - level_warning: typing.ClassVar[ - Logger.LogLevel - ] # value = - def __eq__(self, other: typing.Any) -> bool: ... - def __getstate__(self) -> int: ... - def __hash__(self) -> int: ... - def __index__(self) -> int: ... + level_warning: typing.ClassVar[Logger.LogLevel] # value = + def __eq__(self, other: typing.Any) -> bool: + ... + def __getstate__(self) -> int: + ... + def __hash__(self) -> int: + ... + def __index__(self) -> int: + ... @typing.overload - def __init__(self, value: typing.SupportsInt) -> None: ... + def __init__(self, value: typing.SupportsInt) -> None: + ... @typing.overload - def __init__(self, name: str) -> None: ... - def __int__(self) -> int: ... - def __ne__(self, other: typing.Any) -> bool: ... - def __repr__(self) -> str: ... - def __setstate__(self, state: typing.SupportsInt) -> None: ... - def __str__(self) -> str: ... + def __init__(self, name: str) -> None: + ... + def __int__(self) -> int: + ... + def __ne__(self, other: typing.Any) -> bool: + ... + def __repr__(self) -> str: + ... + def __setstate__(self, state: typing.SupportsInt) -> None: + ... + def __str__(self) -> str: + ... @property - def name(self) -> str: ... + def name(self) -> str: + ... @property - def value(self) -> int: ... - + def value(self) -> int: + ... level_debug: typing.ClassVar[Logger.LogLevel] # value = level_error: typing.ClassVar[Logger.LogLevel] # value = level_info: typing.ClassVar[Logger.LogLevel] # value = - level_warning: typing.ClassVar[ - Logger.LogLevel - ] # value = + level_warning: typing.ClassVar[Logger.LogLevel] # value = @staticmethod - def debug(message: str) -> None: ... + def debug(message: str) -> None: + ... @staticmethod - def error(message: str) -> None: ... + def error(message: str) -> None: + ... @staticmethod - def info(message: str) -> None: ... + def info(message: str) -> None: + ... @staticmethod - def log(level: Logger.LogLevel, message: str) -> None: ... + def log(level: Logger.LogLevel, message: str) -> None: + ... @staticmethod - def set_log_handler( - func: collections.abc.Callable[[Logger.LogLevel, str], None], - ) -> None: ... + def set_log_handler(func: collections.abc.Callable[[Logger.LogLevel, str], None]) -> None: + ... @staticmethod - def warning(message: str) -> None: ... - + def warning(message: str) -> None: + ... class Luminosity(Mapping): - def __init__( - self, - s_lab: typing.SupportsFloat, - s_hat_min: typing.SupportsFloat, - s_hat_max: typing.SupportsFloat = 0.0, - invariant_power: typing.SupportsFloat = 0.0, - mass: typing.SupportsFloat = 0.0, - width: typing.SupportsFloat = 0.0, - ) -> None: ... - + def __init__(self, s_lab: typing.SupportsFloat, s_hat_min: typing.SupportsFloat, s_hat_max: typing.SupportsFloat = 0.0, invariant_power: typing.SupportsFloat = 0.0, mass: typing.SupportsFloat = 0.0, width: typing.SupportsFloat = 0.0) -> None: + ... class MLP(FunctionGenerator): class Activation: """ Members: - + relu - + leaky_relu - + elu - + gelu - + sigmoid - + softplus - + linear """ - - __members__: typing.ClassVar[ - dict[str, MLP.Activation] - ] # value = {'relu': , 'leaky_relu': , 'elu': , 'gelu': , 'sigmoid': , 'softplus': , 'linear': } + __members__: typing.ClassVar[dict[str, MLP.Activation]] # value = {'relu': , 'leaky_relu': , 'elu': , 'gelu': , 'sigmoid': , 'softplus': , 'linear': } elu: typing.ClassVar[MLP.Activation] # value = gelu: typing.ClassVar[MLP.Activation] # value = - leaky_relu: typing.ClassVar[ - MLP.Activation - ] # value = + leaky_relu: typing.ClassVar[MLP.Activation] # value = linear: typing.ClassVar[MLP.Activation] # value = relu: typing.ClassVar[MLP.Activation] # value = sigmoid: typing.ClassVar[MLP.Activation] # value = softplus: typing.ClassVar[MLP.Activation] # value = - def __eq__(self, other: typing.Any) -> bool: ... - def __getstate__(self) -> int: ... - def __hash__(self) -> int: ... - def __index__(self) -> int: ... + def __eq__(self, other: typing.Any) -> bool: + ... + def __getstate__(self) -> int: + ... + def __hash__(self) -> int: + ... + def __index__(self) -> int: + ... @typing.overload - def __init__(self, value: typing.SupportsInt) -> None: ... + def __init__(self, value: typing.SupportsInt) -> None: + ... @typing.overload - def __init__(self, name: str) -> None: ... - def __int__(self) -> int: ... - def __ne__(self, other: typing.Any) -> bool: ... - def __repr__(self) -> str: ... - def __setstate__(self, state: typing.SupportsInt) -> None: ... - def __str__(self) -> str: ... + def __init__(self, name: str) -> None: + ... + def __int__(self) -> int: + ... + def __ne__(self, other: typing.Any) -> bool: + ... + def __repr__(self) -> str: + ... + def __setstate__(self, state: typing.SupportsInt) -> None: + ... + def __str__(self) -> str: + ... @property - def name(self) -> str: ... + def name(self) -> str: + ... @property - def value(self) -> int: ... - + def value(self) -> int: + ... elu: typing.ClassVar[MLP.Activation] # value = gelu: typing.ClassVar[MLP.Activation] # value = leaky_relu: typing.ClassVar[MLP.Activation] # value = @@ -1526,881 +1489,926 @@ class MLP(FunctionGenerator): relu: typing.ClassVar[MLP.Activation] # value = sigmoid: typing.ClassVar[MLP.Activation] # value = softplus: typing.ClassVar[MLP.Activation] # value = - def __init__( - self, - input_dim: typing.SupportsInt, - output_dim: typing.SupportsInt, - hidden_dim: typing.SupportsInt = 32, - layers: typing.SupportsInt = 3, - activation: MLP.Activation = MLP.Activation.Activation.leaky_relu, - prefix: str = "", - ) -> None: ... - def initialize_globals(self, context: Context) -> None: ... - def input_dim(self) -> int: ... - def output_dim(self) -> int: ... - + def __init__(self, input_dim: typing.SupportsInt, output_dim: typing.SupportsInt, hidden_dim: typing.SupportsInt = 32, layers: typing.SupportsInt = 3, activation: MLP.Activation = MLP.Activation.Activation.leaky_relu, prefix: str = '') -> None: + ... + def initialize_globals(self, context: Context) -> None: + ... + def input_dim(self) -> int: + ... + def output_dim(self) -> int: + ... +class MadnisConfig: + lr_schedule: AdamOptimizer.LRSchedule + verbosity: Verbosity + def __init__(self) -> None: + ... + @property + def adam_beta1(self) -> float: + ... + @adam_beta1.setter + def adam_beta1(self, arg0: typing.SupportsFloat) -> None: + ... + @property + def adam_beta2(self) -> float: + ... + @adam_beta2.setter + def adam_beta2(self, arg0: typing.SupportsFloat) -> None: + ... + @property + def adam_eps(self) -> float: + ... + @adam_eps.setter + def adam_eps(self, arg0: typing.SupportsFloat) -> None: + ... + @property + def batch_size_offset(self) -> int: + ... + @batch_size_offset.setter + def batch_size_offset(self, arg0: typing.SupportsInt) -> None: + ... + @property + def batch_size_per_channel(self) -> int: + ... + @batch_size_per_channel.setter + def batch_size_per_channel(self, arg0: typing.SupportsInt) -> None: + ... + @property + def batches(self) -> int: + ... + @batches.setter + def batches(self, arg0: typing.SupportsInt) -> None: + ... + @property + def buffer_capacity(self) -> int: + ... + @buffer_capacity.setter + def buffer_capacity(self, arg0: typing.SupportsInt) -> None: + ... + @property + def buffer_unweighting_quantile(self) -> float: + ... + @buffer_unweighting_quantile.setter + def buffer_unweighting_quantile(self, arg0: typing.SupportsFloat) -> None: + ... + @property + def buffered_steps(self) -> int: + ... + @buffered_steps.setter + def buffered_steps(self, arg0: typing.SupportsInt) -> None: + ... + @property + def channel_dropping_interval(self) -> int: + ... + @channel_dropping_interval.setter + def channel_dropping_interval(self, arg0: typing.SupportsInt) -> None: + ... + @property + def channel_dropping_threshold(self) -> float: + ... + @channel_dropping_threshold.setter + def channel_dropping_threshold(self, arg0: typing.SupportsFloat) -> None: + ... + @property + def cpu_generator_batch_size(self) -> int: + ... + @cpu_generator_batch_size.setter + def cpu_generator_batch_size(self, arg0: typing.SupportsInt) -> None: + ... + @property + def fixed_cwnet_fraction(self) -> float: + ... + @fixed_cwnet_fraction.setter + def fixed_cwnet_fraction(self, arg0: typing.SupportsFloat) -> None: + ... + @property + def generator_target_size_factor(self) -> int: + ... + @generator_target_size_factor.setter + def generator_target_size_factor(self, arg0: typing.SupportsInt) -> None: + ... + @property + def gpu_generator_batch_granularity(self) -> int: + ... + @gpu_generator_batch_granularity.setter + def gpu_generator_batch_granularity(self, arg0: typing.SupportsInt) -> None: + ... + @property + def gpu_generator_batch_size(self) -> int: + ... + @gpu_generator_batch_size.setter + def gpu_generator_batch_size(self, arg0: typing.SupportsInt) -> None: + ... + @property + def integration_history_length(self) -> int: + ... + @integration_history_length.setter + def integration_history_length(self, arg0: typing.SupportsInt) -> None: + ... + @property + def learning_rate(self) -> float: + ... + @learning_rate.setter + def learning_rate(self, arg0: typing.SupportsFloat) -> None: + ... + @property + def log_interval(self) -> int: + ... + @log_interval.setter + def log_interval(self, arg0: typing.SupportsInt) -> None: + ... + @property + def minimum_buffer_size(self) -> int: + ... + @minimum_buffer_size.setter + def minimum_buffer_size(self, arg0: typing.SupportsInt) -> None: + ... + @property + def softclip_threshold(self) -> float: + ... + @softclip_threshold.setter + def softclip_threshold(self, arg0: typing.SupportsFloat) -> None: + ... + @property + def uniform_channel_ratio(self) -> float: + ... + @uniform_channel_ratio.setter + def uniform_channel_ratio(self, arg0: typing.SupportsFloat) -> None: + ... +class MadnisLoss(FunctionGenerator): + def __init__(self, functions: collections.abc.Sequence[FunctionGenerator], cwnet: madspace._madspace_py.ChannelWeightNetwork | None, softclip_threshold: typing.SupportsFloat = 0.0) -> None: + ... +class MadnisTraining: + def __init__(self, generator_context: Context, optimizer_context: Context, config: MadnisConfig, integrands: collections.abc.Sequence[Integrand], cwnet: madspace._madspace_py.ChannelWeightNetwork | None) -> None: + ... + def active_channel_count(self) -> int: + ... + def active_channels(self) -> list[int]: + ... + def train_step(self, batch_index: typing.SupportsInt) -> None: + ... class Mapping: - def __init__( - self, - name: str, - input_types: collections.abc.Sequence[Type], - output_types: collections.abc.Sequence[Type], - condition_types: collections.abc.Sequence[Type], - ) -> None: ... - def build_forward( - self, - builder: FunctionBuilder, - inputs: collections.abc.Sequence[Value], - conditions: collections.abc.Sequence[Value], - ) -> tuple[list[Value], Value]: ... - def build_inverse( - self, - builder: FunctionBuilder, - inputs: collections.abc.Sequence[Value], - conditions: collections.abc.Sequence[Value], - ) -> tuple[list[Value], Value]: ... - def forward_function(self) -> Function: ... - def inverse_function(self) -> Function: ... - def map_forward(self, inputs, conditions=list()): ... - def map_inverse(self, inputs, conditions=list()): ... - + def __init__(self, name: str, input_types: NamedTypes, output_types: NamedTypes, condition_types: NamedTypes) -> None: + ... + @typing.overload + def build_forward(self, builder: FunctionBuilder, inputs: collections.abc.Sequence[Value], conditions: collections.abc.Sequence[Value]) -> NamedValues: + ... + @typing.overload + def build_forward(self, builder: FunctionBuilder, inputs: NamedValues, conditions: NamedValues) -> NamedValues: + ... + @typing.overload + def build_inverse(self, builder: FunctionBuilder, inputs: collections.abc.Sequence[Value], conditions: collections.abc.Sequence[Value]) -> NamedValues: + ... + @typing.overload + def build_inverse(self, builder: FunctionBuilder, inputs: NamedValues, conditions: NamedValues) -> NamedValues: + ... + def forward_function(self) -> Function: + ... + def inverse_function(self) -> Function: + ... + def map_forward(self, inputs = list(), conditions = list(), **kwargs): + ... + def map_inverse(self, inputs = list(), conditions = list(), **kwargs): + ... class MatrixElement(FunctionGenerator): class MatrixElementInput: """ Members: - + momenta_in - + alpha_s_in - + flavor_in - + random_color_in - + random_helicity_in - + random_diagram_in - + helicity_in - + diagram_in """ - - __members__: typing.ClassVar[ - dict[str, MatrixElement.MatrixElementInput] - ] # value = {'momenta_in': , 'alpha_s_in': , 'flavor_in': , 'random_color_in': , 'random_helicity_in': , 'random_diagram_in': , 'helicity_in': , 'diagram_in': } - alpha_s_in: typing.ClassVar[ - MatrixElement.MatrixElementInput - ] # value = - diagram_in: typing.ClassVar[ - MatrixElement.MatrixElementInput - ] # value = - flavor_in: typing.ClassVar[ - MatrixElement.MatrixElementInput - ] # value = - helicity_in: typing.ClassVar[ - MatrixElement.MatrixElementInput - ] # value = - momenta_in: typing.ClassVar[ - MatrixElement.MatrixElementInput - ] # value = - random_color_in: typing.ClassVar[ - MatrixElement.MatrixElementInput - ] # value = - random_diagram_in: typing.ClassVar[ - MatrixElement.MatrixElementInput - ] # value = - random_helicity_in: typing.ClassVar[ - MatrixElement.MatrixElementInput - ] # value = - def __eq__(self, other: typing.Any) -> bool: ... - def __getstate__(self) -> int: ... - def __hash__(self) -> int: ... - def __index__(self) -> int: ... + __members__: typing.ClassVar[dict[str, MatrixElement.MatrixElementInput]] # value = {'momenta_in': , 'alpha_s_in': , 'flavor_in': , 'random_color_in': , 'random_helicity_in': , 'random_diagram_in': , 'helicity_in': , 'diagram_in': } + alpha_s_in: typing.ClassVar[MatrixElement.MatrixElementInput] # value = + diagram_in: typing.ClassVar[MatrixElement.MatrixElementInput] # value = + flavor_in: typing.ClassVar[MatrixElement.MatrixElementInput] # value = + helicity_in: typing.ClassVar[MatrixElement.MatrixElementInput] # value = + momenta_in: typing.ClassVar[MatrixElement.MatrixElementInput] # value = + random_color_in: typing.ClassVar[MatrixElement.MatrixElementInput] # value = + random_diagram_in: typing.ClassVar[MatrixElement.MatrixElementInput] # value = + random_helicity_in: typing.ClassVar[MatrixElement.MatrixElementInput] # value = + def __eq__(self, other: typing.Any) -> bool: + ... + def __getstate__(self) -> int: + ... + def __hash__(self) -> int: + ... + def __index__(self) -> int: + ... @typing.overload - def __init__(self, value: typing.SupportsInt) -> None: ... + def __init__(self, value: typing.SupportsInt) -> None: + ... @typing.overload - def __init__(self, name: str) -> None: ... - def __int__(self) -> int: ... - def __ne__(self, other: typing.Any) -> bool: ... - def __repr__(self) -> str: ... - def __setstate__(self, state: typing.SupportsInt) -> None: ... - def __str__(self) -> str: ... + def __init__(self, name: str) -> None: + ... + def __int__(self) -> int: + ... + def __ne__(self, other: typing.Any) -> bool: + ... + def __repr__(self) -> str: + ... + def __setstate__(self, state: typing.SupportsInt) -> None: + ... + def __str__(self) -> str: + ... @property - def name(self) -> str: ... + def name(self) -> str: + ... @property - def value(self) -> int: ... - + def value(self) -> int: + ... class MatrixElementOutput: """ Members: - + matrix_element_out - + diagram_amp2_out - + color_index_out - + helicity_index_out - + diagram_index_out """ - - __members__: typing.ClassVar[ - dict[str, MatrixElement.MatrixElementOutput] - ] # value = {'matrix_element_out': , 'diagram_amp2_out': , 'color_index_out': , 'helicity_index_out': , 'diagram_index_out': } - color_index_out: typing.ClassVar[ - MatrixElement.MatrixElementOutput - ] # value = - diagram_amp2_out: typing.ClassVar[ - MatrixElement.MatrixElementOutput - ] # value = - diagram_index_out: typing.ClassVar[ - MatrixElement.MatrixElementOutput - ] # value = - helicity_index_out: typing.ClassVar[ - MatrixElement.MatrixElementOutput - ] # value = - matrix_element_out: typing.ClassVar[ - MatrixElement.MatrixElementOutput - ] # value = - def __eq__(self, other: typing.Any) -> bool: ... - def __getstate__(self) -> int: ... - def __hash__(self) -> int: ... - def __index__(self) -> int: ... + __members__: typing.ClassVar[dict[str, MatrixElement.MatrixElementOutput]] # value = {'matrix_element_out': , 'diagram_amp2_out': , 'color_index_out': , 'helicity_index_out': , 'diagram_index_out': } + color_index_out: typing.ClassVar[MatrixElement.MatrixElementOutput] # value = + diagram_amp2_out: typing.ClassVar[MatrixElement.MatrixElementOutput] # value = + diagram_index_out: typing.ClassVar[MatrixElement.MatrixElementOutput] # value = + helicity_index_out: typing.ClassVar[MatrixElement.MatrixElementOutput] # value = + matrix_element_out: typing.ClassVar[MatrixElement.MatrixElementOutput] # value = + def __eq__(self, other: typing.Any) -> bool: + ... + def __getstate__(self) -> int: + ... + def __hash__(self) -> int: + ... + def __index__(self) -> int: + ... @typing.overload - def __init__(self, value: typing.SupportsInt) -> None: ... + def __init__(self, value: typing.SupportsInt) -> None: + ... @typing.overload - def __init__(self, name: str) -> None: ... - def __int__(self) -> int: ... - def __ne__(self, other: typing.Any) -> bool: ... - def __repr__(self) -> str: ... - def __setstate__(self, state: typing.SupportsInt) -> None: ... - def __str__(self) -> str: ... + def __init__(self, name: str) -> None: + ... + def __int__(self) -> int: + ... + def __ne__(self, other: typing.Any) -> bool: + ... + def __repr__(self) -> str: + ... + def __setstate__(self, state: typing.SupportsInt) -> None: + ... + def __str__(self) -> str: + ... @property - def name(self) -> str: ... + def name(self) -> str: + ... @property - def value(self) -> int: ... - - alpha_s_in: typing.ClassVar[ - MatrixElement.MatrixElementInput - ] # value = - color_index_out: typing.ClassVar[ - MatrixElement.MatrixElementOutput - ] # value = - diagram_amp2_out: typing.ClassVar[ - MatrixElement.MatrixElementOutput - ] # value = - diagram_in: typing.ClassVar[ - MatrixElement.MatrixElementInput - ] # value = - diagram_index_out: typing.ClassVar[ - MatrixElement.MatrixElementOutput - ] # value = - flavor_in: typing.ClassVar[ - MatrixElement.MatrixElementInput - ] # value = - helicity_in: typing.ClassVar[ - MatrixElement.MatrixElementInput - ] # value = - helicity_index_out: typing.ClassVar[ - MatrixElement.MatrixElementOutput - ] # value = - matrix_element_out: typing.ClassVar[ - MatrixElement.MatrixElementOutput - ] # value = - momenta_in: typing.ClassVar[ - MatrixElement.MatrixElementInput - ] # value = - random_color_in: typing.ClassVar[ - MatrixElement.MatrixElementInput - ] # value = - random_diagram_in: typing.ClassVar[ - MatrixElement.MatrixElementInput - ] # value = - random_helicity_in: typing.ClassVar[ - MatrixElement.MatrixElementInput - ] # value = + def value(self) -> int: + ... + alpha_s_in: typing.ClassVar[MatrixElement.MatrixElementInput] # value = + color_index_out: typing.ClassVar[MatrixElement.MatrixElementOutput] # value = + diagram_amp2_out: typing.ClassVar[MatrixElement.MatrixElementOutput] # value = + diagram_in: typing.ClassVar[MatrixElement.MatrixElementInput] # value = + diagram_index_out: typing.ClassVar[MatrixElement.MatrixElementOutput] # value = + flavor_in: typing.ClassVar[MatrixElement.MatrixElementInput] # value = + helicity_in: typing.ClassVar[MatrixElement.MatrixElementInput] # value = + helicity_index_out: typing.ClassVar[MatrixElement.MatrixElementOutput] # value = + matrix_element_out: typing.ClassVar[MatrixElement.MatrixElementOutput] # value = + momenta_in: typing.ClassVar[MatrixElement.MatrixElementInput] # value = + random_color_in: typing.ClassVar[MatrixElement.MatrixElementInput] # value = + random_diagram_in: typing.ClassVar[MatrixElement.MatrixElementInput] # value = + random_helicity_in: typing.ClassVar[MatrixElement.MatrixElementInput] # value = @typing.overload - def __init__( - self, - matrix_element_index: typing.SupportsInt, - particle_count: typing.SupportsInt, - inputs: collections.abc.Sequence[MatrixElement.MatrixElementInput], - outputs: collections.abc.Sequence[MatrixElement.MatrixElementOutput], - diagram_count: typing.SupportsInt = 1, - sample_random_inputs: bool = False, - ) -> None: ... + def __init__(self, matrix_element_index: typing.SupportsInt, particle_count: typing.SupportsInt, inputs: collections.abc.Sequence[MatrixElement.MatrixElementInput], outputs: collections.abc.Sequence[MatrixElement.MatrixElementOutput], diagram_count: typing.SupportsInt = 1, sample_random_inputs: bool = False) -> None: + ... @typing.overload - def __init__( - self, - matrix_element_api: MatrixElementApi, - inputs: collections.abc.Sequence[MatrixElement.MatrixElementInput], - outputs: collections.abc.Sequence[MatrixElement.MatrixElementOutput], - sample_random_inputs: bool = False, - ) -> None: ... - def diagram_count(self) -> int: ... - def matrix_element_index(self) -> int: ... - def particle_count(self) -> int: ... - + def __init__(self, matrix_element_api: MatrixElementApi, inputs: collections.abc.Sequence[MatrixElement.MatrixElementInput], outputs: collections.abc.Sequence[MatrixElement.MatrixElementOutput], sample_random_inputs: bool = False) -> None: + ... + def diagram_count(self) -> int: + ... + def matrix_element_index(self) -> int: + ... + def particle_count(self) -> int: + ... class MatrixElementApi: - def __init__( - self, file: str, param_card: str, index: typing.SupportsInt = 0 - ) -> None: ... - def device(self) -> Device: ... - def diagram_count(self) -> int: ... - def helicity_count(self) -> int: ... - def index(self) -> int: ... - def particle_count(self) -> int: ... - + def diagram_count(self) -> int: + ... + def helicity_count(self) -> int: + ... + def index(self) -> int: + ... + def particle_count(self) -> int: + ... class MomentumPreprocessing(FunctionGenerator): - def __init__(self, particle_count: typing.SupportsInt) -> None: ... - def output_dim(self) -> int: ... - + def __init__(self, particle_count: typing.SupportsInt) -> None: + ... + def output_dim(self) -> int: + ... class MultiChannelFunction(FunctionGenerator): - def __init__( - self, functions: collections.abc.Sequence[FunctionGenerator] - ) -> None: ... - + def __init__(self, functions: collections.abc.Sequence[FunctionGenerator]) -> None: + ... class MultiChannelIntegrand(FunctionGenerator): - def __init__(self, integrands: collections.abc.Sequence[Integrand]) -> None: ... - + def __init__(self, integrands: collections.abc.Sequence[Integrand], return_sizes: bool = False) -> None: + ... class MultiChannelMapping(Mapping): - def __init__(self, mappings: collections.abc.Sequence[Mapping]) -> None: ... - + def __init__(self, mappings: collections.abc.Sequence[Mapping]) -> None: + ... +class MultiMadnisTraining: + def __init__(self, generator_context: Context, optimizer_context: Context, config: MadnisConfig, integrands: collections.abc.Sequence[collections.abc.Sequence[Integrand]], cwnets: collections.abc.Sequence[madspace._madspace_py.ChannelWeightNetwork | None]) -> None: + ... + def active_channels(self) -> list[list[int]]: + ... + def train(self) -> None: + ... +class NamedTypes: + @typing.overload + def __getitem__(self, key: typing.SupportsInt) -> Type: + ... + @typing.overload + def __getitem__(self, key: str) -> Type: + ... + @typing.overload + def __init__(self) -> None: + ... + @typing.overload + def __init__(self, keys: collections.abc.Sequence[str], values: collections.abc.Sequence[Type]) -> None: + ... + @typing.overload + def __init__(self, items: collections.abc.Sequence[tuple[str, Type]]) -> None: + ... + def __len__(self) -> int: + ... + def index_map(self) -> dict[str, int]: + ... + def keys(self) -> list[str]: + ... + def push_back(self, name: str, item: Type) -> None: + ... + def values(self) -> list[Type]: + ... +class NamedValues: + @typing.overload + def __getitem__(self, key: typing.SupportsInt) -> Value: + ... + @typing.overload + def __getitem__(self, key: str) -> Value: + ... + @typing.overload + def __init__(self) -> None: + ... + @typing.overload + def __init__(self, keys: collections.abc.Sequence[str], values: collections.abc.Sequence[Value]) -> None: + ... + @typing.overload + def __init__(self, items: collections.abc.Sequence[tuple[str, Value]]) -> None: + ... + def __len__(self) -> int: + ... + def index_map(self) -> dict[str, int]: + ... + def keys(self) -> list[str]: + ... + def push_back(self, name: str, item: Value) -> None: + ... + def values(self) -> list[Value]: + ... class Observable(FunctionGenerator): class ObservableOption: """ Members: - + obs_e - + obs_px - + obs_py - + obs_pz - + obs_mass - + obs_pt - + obs_p_mag - + obs_phi - + obs_theta - + obs_y - + obs_y_abs - + obs_eta - + obs_eta_abs - + obs_delta_eta - + obs_delta_phi - + obs_delta_r - + obs_sqrt_s """ - - __members__: typing.ClassVar[ - dict[str, Observable.ObservableOption] - ] # value = {'obs_e': , 'obs_px': , 'obs_py': , 'obs_pz': , 'obs_mass': , 'obs_pt': , 'obs_p_mag': , 'obs_phi': , 'obs_theta': , 'obs_y': , 'obs_y_abs': , 'obs_eta': , 'obs_eta_abs': , 'obs_delta_eta': , 'obs_delta_phi': , 'obs_delta_r': , 'obs_sqrt_s': } - obs_delta_eta: typing.ClassVar[ - Observable.ObservableOption - ] # value = - obs_delta_phi: typing.ClassVar[ - Observable.ObservableOption - ] # value = - obs_delta_r: typing.ClassVar[ - Observable.ObservableOption - ] # value = - obs_e: typing.ClassVar[ - Observable.ObservableOption - ] # value = - obs_eta: typing.ClassVar[ - Observable.ObservableOption - ] # value = - obs_eta_abs: typing.ClassVar[ - Observable.ObservableOption - ] # value = - obs_mass: typing.ClassVar[ - Observable.ObservableOption - ] # value = - obs_p_mag: typing.ClassVar[ - Observable.ObservableOption - ] # value = - obs_phi: typing.ClassVar[ - Observable.ObservableOption - ] # value = - obs_pt: typing.ClassVar[ - Observable.ObservableOption - ] # value = - obs_px: typing.ClassVar[ - Observable.ObservableOption - ] # value = - obs_py: typing.ClassVar[ - Observable.ObservableOption - ] # value = - obs_pz: typing.ClassVar[ - Observable.ObservableOption - ] # value = - obs_sqrt_s: typing.ClassVar[ - Observable.ObservableOption - ] # value = - obs_theta: typing.ClassVar[ - Observable.ObservableOption - ] # value = - obs_y: typing.ClassVar[ - Observable.ObservableOption - ] # value = - obs_y_abs: typing.ClassVar[ - Observable.ObservableOption - ] # value = - def __eq__(self, other: typing.Any) -> bool: ... - def __getstate__(self) -> int: ... - def __hash__(self) -> int: ... - def __index__(self) -> int: ... + __members__: typing.ClassVar[dict[str, Observable.ObservableOption]] # value = {'obs_e': , 'obs_px': , 'obs_py': , 'obs_pz': , 'obs_mass': , 'obs_pt': , 'obs_p_mag': , 'obs_phi': , 'obs_theta': , 'obs_y': , 'obs_y_abs': , 'obs_eta': , 'obs_eta_abs': , 'obs_delta_eta': , 'obs_delta_phi': , 'obs_delta_r': , 'obs_sqrt_s': } + obs_delta_eta: typing.ClassVar[Observable.ObservableOption] # value = + obs_delta_phi: typing.ClassVar[Observable.ObservableOption] # value = + obs_delta_r: typing.ClassVar[Observable.ObservableOption] # value = + obs_e: typing.ClassVar[Observable.ObservableOption] # value = + obs_eta: typing.ClassVar[Observable.ObservableOption] # value = + obs_eta_abs: typing.ClassVar[Observable.ObservableOption] # value = + obs_mass: typing.ClassVar[Observable.ObservableOption] # value = + obs_p_mag: typing.ClassVar[Observable.ObservableOption] # value = + obs_phi: typing.ClassVar[Observable.ObservableOption] # value = + obs_pt: typing.ClassVar[Observable.ObservableOption] # value = + obs_px: typing.ClassVar[Observable.ObservableOption] # value = + obs_py: typing.ClassVar[Observable.ObservableOption] # value = + obs_pz: typing.ClassVar[Observable.ObservableOption] # value = + obs_sqrt_s: typing.ClassVar[Observable.ObservableOption] # value = + obs_theta: typing.ClassVar[Observable.ObservableOption] # value = + obs_y: typing.ClassVar[Observable.ObservableOption] # value = + obs_y_abs: typing.ClassVar[Observable.ObservableOption] # value = + def __eq__(self, other: typing.Any) -> bool: + ... + def __getstate__(self) -> int: + ... + def __hash__(self) -> int: + ... + def __index__(self) -> int: + ... @typing.overload - def __init__(self, value: typing.SupportsInt) -> None: ... + def __init__(self, value: typing.SupportsInt) -> None: + ... @typing.overload - def __init__(self, name: str) -> None: ... - def __int__(self) -> int: ... - def __ne__(self, other: typing.Any) -> bool: ... - def __repr__(self) -> str: ... - def __setstate__(self, state: typing.SupportsInt) -> None: ... - def __str__(self) -> str: ... + def __init__(self, name: str) -> None: + ... + def __int__(self) -> int: + ... + def __ne__(self, other: typing.Any) -> bool: + ... + def __repr__(self) -> str: + ... + def __setstate__(self, state: typing.SupportsInt) -> None: + ... + def __str__(self) -> str: + ... @property - def name(self) -> str: ... + def name(self) -> str: + ... @property - def value(self) -> int: ... - + def value(self) -> int: + ... bottom_pids: typing.ClassVar[list] = [-5, 5] jet_pids: typing.ClassVar[list] = [1, 2, 3, 4, -1, -2, -3, -4, 21] lepton_pids: typing.ClassVar[list] = [11, 13, 15, -11, -13, -15] missing_pids: typing.ClassVar[list] = [12, 14, 16, -12, -14, -16] - obs_delta_eta: typing.ClassVar[ - Observable.ObservableOption - ] # value = - obs_delta_phi: typing.ClassVar[ - Observable.ObservableOption - ] # value = - obs_delta_r: typing.ClassVar[ - Observable.ObservableOption - ] # value = - obs_e: typing.ClassVar[ - Observable.ObservableOption - ] # value = - obs_eta: typing.ClassVar[ - Observable.ObservableOption - ] # value = - obs_eta_abs: typing.ClassVar[ - Observable.ObservableOption - ] # value = - obs_mass: typing.ClassVar[ - Observable.ObservableOption - ] # value = - obs_p_mag: typing.ClassVar[ - Observable.ObservableOption - ] # value = - obs_phi: typing.ClassVar[ - Observable.ObservableOption - ] # value = - obs_pt: typing.ClassVar[ - Observable.ObservableOption - ] # value = - obs_px: typing.ClassVar[ - Observable.ObservableOption - ] # value = - obs_py: typing.ClassVar[ - Observable.ObservableOption - ] # value = - obs_pz: typing.ClassVar[ - Observable.ObservableOption - ] # value = - obs_sqrt_s: typing.ClassVar[ - Observable.ObservableOption - ] # value = - obs_theta: typing.ClassVar[ - Observable.ObservableOption - ] # value = - obs_y: typing.ClassVar[ - Observable.ObservableOption - ] # value = - obs_y_abs: typing.ClassVar[ - Observable.ObservableOption - ] # value = + obs_delta_eta: typing.ClassVar[Observable.ObservableOption] # value = + obs_delta_phi: typing.ClassVar[Observable.ObservableOption] # value = + obs_delta_r: typing.ClassVar[Observable.ObservableOption] # value = + obs_e: typing.ClassVar[Observable.ObservableOption] # value = + obs_eta: typing.ClassVar[Observable.ObservableOption] # value = + obs_eta_abs: typing.ClassVar[Observable.ObservableOption] # value = + obs_mass: typing.ClassVar[Observable.ObservableOption] # value = + obs_p_mag: typing.ClassVar[Observable.ObservableOption] # value = + obs_phi: typing.ClassVar[Observable.ObservableOption] # value = + obs_pt: typing.ClassVar[Observable.ObservableOption] # value = + obs_px: typing.ClassVar[Observable.ObservableOption] # value = + obs_py: typing.ClassVar[Observable.ObservableOption] # value = + obs_pz: typing.ClassVar[Observable.ObservableOption] # value = + obs_sqrt_s: typing.ClassVar[Observable.ObservableOption] # value = + obs_theta: typing.ClassVar[Observable.ObservableOption] # value = + obs_y: typing.ClassVar[Observable.ObservableOption] # value = + obs_y_abs: typing.ClassVar[Observable.ObservableOption] # value = photon_pids: typing.ClassVar[list] = [22] - def __init__( - self, - pids: collections.abc.Sequence[typing.SupportsInt], - observable: Observable.ObservableOption, - select_pids: collections.abc.Sequence[ - collections.abc.Sequence[typing.SupportsInt] - ], - sum_momenta: bool = False, - sum_observable: bool = False, - order_observable: ( - madspace._madspace_py.Observable.ObservableOption | None - ) = None, - order_indices: collections.abc.Sequence[typing.SupportsInt] = [], - ignore_incoming: bool = True, - name: str = "", - ) -> None: ... - + def __init__(self, pids: collections.abc.Sequence[typing.SupportsInt], observable: Observable.ObservableOption, select_pids: collections.abc.Sequence[collections.abc.Sequence[typing.SupportsInt]], sum_momenta: bool = False, sum_observable: bool = False, order_observable: madspace._madspace_py.Observable.ObservableOption | None = None, order_indices: collections.abc.Sequence[typing.SupportsInt] = [], ignore_incoming: bool = True, name: str = '') -> None: + ... class ObservableHistograms(FunctionGenerator): - def __init__(self, observables: collections.abc.Sequence[HistItem]) -> None: ... - + def __init__(self, observables: collections.abc.Sequence[HistItem]) -> None: + ... class PartonDensity(FunctionGenerator): - def __init__( - self, - grid: PdfGrid, - pids: collections.abc.Sequence[typing.SupportsInt], - dynamic_pid: bool = False, - prefix: str = "", - ) -> None: ... - + def __init__(self, grid: PdfGrid, pids: collections.abc.Sequence[typing.SupportsInt], dynamic_pid: bool = False, prefix: str = '') -> None: + ... class PdfGrid: - def __init__(self, file: str) -> None: ... - def coefficients_shape(self, batch_dim: bool = False) -> list[int]: ... - def initialize_globals(self, context: Context, prefix: str = "") -> None: ... - def logq2_shape(self, batch_dim: bool = False) -> list[int]: ... - def logx_shape(self, batch_dim: bool = False) -> list[int]: ... + def __init__(self, file: str) -> None: + ... + def coefficients_shape(self, batch_dim: bool = False) -> list[int]: + ... + def initialize_globals(self, context: Context, prefix: str = '') -> None: + ... + def logq2_shape(self, batch_dim: bool = False) -> list[int]: + ... + def logx_shape(self, batch_dim: bool = False) -> list[int]: + ... @property - def grid_point_count(self) -> int: ... + def grid_point_count(self) -> int: + ... @property - def logq2(self) -> list[float]: ... + def logq2(self) -> list[float]: + ... @property - def logx(self) -> list[float]: ... + def logx(self) -> list[float]: + ... @property - def pids(self) -> list[int]: ... + def pids(self) -> list[int]: + ... @property - def q(self) -> list[float]: ... + def q(self) -> list[float]: + ... @property - def q_count(self) -> int: ... + def q_count(self) -> int: + ... @property - def region_sizes(self) -> list[int]: ... + def region_sizes(self) -> list[int]: + ... @property - def values(self) -> list[list[float]]: ... + def values(self) -> list[list[float]]: + ... @property - def x(self) -> list[float]: ... - + def x(self) -> list[float]: + ... class PhaseSpaceMapping(Mapping): class TChannelMode: """ Members: - + propagator - + rambo - + chili """ - - __members__: typing.ClassVar[ - dict[str, PhaseSpaceMapping.TChannelMode] - ] # value = {'propagator': , 'rambo': , 'chili': } - chili: typing.ClassVar[ - PhaseSpaceMapping.TChannelMode - ] # value = - propagator: typing.ClassVar[ - PhaseSpaceMapping.TChannelMode - ] # value = - rambo: typing.ClassVar[ - PhaseSpaceMapping.TChannelMode - ] # value = - def __eq__(self, other: typing.Any) -> bool: ... - def __getstate__(self) -> int: ... - def __hash__(self) -> int: ... - def __index__(self) -> int: ... + __members__: typing.ClassVar[dict[str, PhaseSpaceMapping.TChannelMode]] # value = {'propagator': , 'rambo': , 'chili': } + chili: typing.ClassVar[PhaseSpaceMapping.TChannelMode] # value = + propagator: typing.ClassVar[PhaseSpaceMapping.TChannelMode] # value = + rambo: typing.ClassVar[PhaseSpaceMapping.TChannelMode] # value = + def __eq__(self, other: typing.Any) -> bool: + ... + def __getstate__(self) -> int: + ... + def __hash__(self) -> int: + ... + def __index__(self) -> int: + ... @typing.overload - def __init__(self, value: typing.SupportsInt) -> None: ... + def __init__(self, value: typing.SupportsInt) -> None: + ... @typing.overload - def __init__(self, name: str) -> None: ... - def __int__(self) -> int: ... - def __ne__(self, other: typing.Any) -> bool: ... - def __repr__(self) -> str: ... - def __setstate__(self, state: typing.SupportsInt) -> None: ... - def __str__(self) -> str: ... + def __init__(self, name: str) -> None: + ... + def __int__(self) -> int: + ... + def __ne__(self, other: typing.Any) -> bool: + ... + def __repr__(self) -> str: + ... + def __setstate__(self, state: typing.SupportsInt) -> None: + ... + def __str__(self) -> str: + ... @property - def name(self) -> str: ... + def name(self) -> str: + ... @property - def value(self) -> int: ... - - chili: typing.ClassVar[ - PhaseSpaceMapping.TChannelMode - ] # value = - propagator: typing.ClassVar[ - PhaseSpaceMapping.TChannelMode - ] # value = - rambo: typing.ClassVar[ - PhaseSpaceMapping.TChannelMode - ] # value = + def value(self) -> int: + ... + chili: typing.ClassVar[PhaseSpaceMapping.TChannelMode] # value = + propagator: typing.ClassVar[PhaseSpaceMapping.TChannelMode] # value = + rambo: typing.ClassVar[PhaseSpaceMapping.TChannelMode] # value = @typing.overload - def __init__( - self, - topology: Topology, - cm_energy: typing.SupportsFloat, - leptonic: bool = False, - invariant_power: typing.SupportsFloat = 0.8, - t_channel_mode: PhaseSpaceMapping.TChannelMode = PhaseSpaceMapping.TChannelMode.TChannelMode.propagator, - cuts: madspace._madspace_py.Cuts | None = None, - permutations: collections.abc.Sequence[ - collections.abc.Sequence[typing.SupportsInt] - ] = [], - ) -> None: ... + def __init__(self, topology: Topology, cm_energy: typing.SupportsFloat, leptonic: bool = False, invariant_power: typing.SupportsFloat = 0.8, t_channel_mode: PhaseSpaceMapping.TChannelMode = PhaseSpaceMapping.TChannelMode.TChannelMode.propagator, cuts: madspace._madspace_py.Cuts | None = None, permutations: collections.abc.Sequence[collections.abc.Sequence[typing.SupportsInt]] = []) -> None: + ... @typing.overload - def __init__( - self, - masses: collections.abc.Sequence[typing.SupportsFloat], - cm_energy: typing.SupportsFloat, - leptonic: bool = False, - invariant_power: typing.SupportsFloat = 0.8, - mode: PhaseSpaceMapping.TChannelMode = PhaseSpaceMapping.TChannelMode.TChannelMode.rambo, - cuts: madspace._madspace_py.Cuts | None = None, - ) -> None: ... - def channel_count(self) -> int: ... - def particle_count(self) -> int: ... - def random_dim(self) -> int: ... - + def __init__(self, masses: collections.abc.Sequence[typing.SupportsFloat], cm_energy: typing.SupportsFloat, leptonic: bool = False, invariant_power: typing.SupportsFloat = 0.8, mode: PhaseSpaceMapping.TChannelMode = PhaseSpaceMapping.TChannelMode.TChannelMode.rambo, cuts: madspace._madspace_py.Cuts | None = None) -> None: + ... + def channel_count(self) -> int: + ... + def particle_count(self) -> int: + ... + def random_dim(self) -> int: + ... class PrettyBox: - def __init__( - self, - title: str, - rows: typing.SupportsInt, - columns: collections.abc.Sequence[typing.SupportsInt], - offset: typing.SupportsInt = 0, - box_width: typing.SupportsInt = 91, - ) -> None: ... - def print_first(self) -> None: ... - def print_update(self) -> None: ... - def set_cell( - self, row: typing.SupportsInt, column: typing.SupportsInt, value: str - ) -> None: ... - def set_column( - self, column: typing.SupportsInt, values: collections.abc.Sequence[str] - ) -> None: ... - def set_row( - self, row: typing.SupportsInt, values: collections.abc.Sequence[str] - ) -> None: ... - @property - def line_count(self) -> int: ... - + def __init__(self, title: str, rows: typing.SupportsInt, columns: collections.abc.Sequence[typing.SupportsInt], offset: typing.SupportsInt = 0, box_width: typing.SupportsInt = 91) -> None: + ... + def print_first(self) -> None: + ... + def print_update(self) -> None: + ... + def set_cell(self, row: typing.SupportsInt, column: typing.SupportsInt, value: str) -> None: + ... + def set_column(self, column: typing.SupportsInt, values: collections.abc.Sequence[str]) -> None: + ... + def set_row(self, row: typing.SupportsInt, values: collections.abc.Sequence[str]) -> None: + ... + @property + def line_count(self) -> int: + ... class Propagator: - def __init__( - self, - mass: typing.SupportsFloat = 0.0, - width: typing.SupportsFloat = 0.0, - integration_order: typing.SupportsInt = 0, - e_min: typing.SupportsFloat = 0.0, - e_max: typing.SupportsFloat = 0.0, - pdg_id: typing.SupportsInt = 0, - ) -> None: ... + def __init__(self, mass: typing.SupportsFloat = 0.0, width: typing.SupportsFloat = 0.0, integration_order: typing.SupportsInt = 0, e_min: typing.SupportsFloat = 0.0, e_max: typing.SupportsFloat = 0.0, pdg_id: typing.SupportsInt = 0) -> None: + ... @property - def e_max(self) -> float: ... + def e_max(self) -> float: + ... @property - def e_min(self) -> float: ... + def e_min(self) -> float: + ... @property - def integration_order(self) -> int: ... + def integration_order(self) -> int: + ... @property - def mass(self) -> float: ... + def mass(self) -> float: + ... @property - def pdg_id(self) -> int: ... + def pdg_id(self) -> int: + ... @property - def width(self) -> float: ... - + def width(self) -> float: + ... class PropagatorChannelWeights(FunctionGenerator): - def __init__( - self, - topologies: collections.abc.Sequence[Topology], - permutations: collections.abc.Sequence[ - collections.abc.Sequence[collections.abc.Sequence[typing.SupportsInt]] - ], - channel_indices: collections.abc.Sequence[ - collections.abc.Sequence[typing.SupportsInt] - ], - ) -> None: ... - + def __init__(self, topologies: collections.abc.Sequence[Topology], permutations: collections.abc.Sequence[collections.abc.Sequence[collections.abc.Sequence[typing.SupportsInt]]], channel_indices: collections.abc.Sequence[collections.abc.Sequence[typing.SupportsInt]]) -> None: + ... class RunningCoupling(FunctionGenerator): - def __init__(self, grid: AlphaSGrid, prefix: str = "") -> None: ... - + def __init__(self, grid: AlphaSGrid, prefix: str = '') -> None: + ... class SubchannelWeights(FunctionGenerator): - def __init__( - self, - topologies: collections.abc.Sequence[collections.abc.Sequence[Topology]], - permutations: collections.abc.Sequence[ - collections.abc.Sequence[collections.abc.Sequence[typing.SupportsInt]] - ], - channel_indices: collections.abc.Sequence[ - collections.abc.Sequence[typing.SupportsInt] - ], - ) -> None: ... - def channel_count(self) -> int: ... - + def __init__(self, topologies: collections.abc.Sequence[collections.abc.Sequence[Topology]], permutations: collections.abc.Sequence[collections.abc.Sequence[collections.abc.Sequence[typing.SupportsInt]]], channel_indices: collections.abc.Sequence[collections.abc.Sequence[typing.SupportsInt]]) -> None: + ... + def channel_count(self) -> int: + ... class SubprocArgs: - def __init__( - self, - process_id: typing.SupportsInt = 0, - topologies: collections.abc.Sequence[Topology] = [], - permutations: collections.abc.Sequence[ - collections.abc.Sequence[collections.abc.Sequence[typing.SupportsInt]] - ] = [], - diagram_indices: collections.abc.Sequence[ - collections.abc.Sequence[typing.SupportsInt] - ] = [], - diagram_color_indices: collections.abc.Sequence[ - collections.abc.Sequence[collections.abc.Sequence[typing.SupportsInt]] - ] = [], - color_flows: collections.abc.Sequence[ - collections.abc.Sequence[ - collections.abc.Sequence[tuple[typing.SupportsInt, typing.SupportsInt]] - ] - ] = [], - pdg_color_types: collections.abc.Mapping[ - typing.SupportsInt, typing.SupportsInt - ] = {}, - helicities: collections.abc.Sequence[ - collections.abc.Sequence[typing.SupportsFloat] - ] = [], - pdg_ids: collections.abc.Sequence[ - collections.abc.Sequence[collections.abc.Sequence[typing.SupportsInt]] - ] = [], - matrix_flavor_indices: collections.abc.Sequence[typing.SupportsInt] = [], - ) -> None: ... - @property - def color_flows(self) -> list[list[list[tuple[int, int]]]]: ... + def __init__(self, process_id: typing.SupportsInt = 0, topologies: collections.abc.Sequence[Topology] = [], permutations: collections.abc.Sequence[collections.abc.Sequence[collections.abc.Sequence[typing.SupportsInt]]] = [], diagram_indices: collections.abc.Sequence[collections.abc.Sequence[typing.SupportsInt]] = [], diagram_color_indices: collections.abc.Sequence[collections.abc.Sequence[collections.abc.Sequence[typing.SupportsInt]]] = [], color_flows: collections.abc.Sequence[collections.abc.Sequence[collections.abc.Sequence[tuple[typing.SupportsInt, typing.SupportsInt]]]] = [], pdg_color_types: collections.abc.Mapping[typing.SupportsInt, typing.SupportsInt] = {}, helicities: collections.abc.Sequence[collections.abc.Sequence[typing.SupportsFloat]] = [], pdg_ids: collections.abc.Sequence[collections.abc.Sequence[collections.abc.Sequence[typing.SupportsInt]]] = [], matrix_flavor_indices: collections.abc.Sequence[typing.SupportsInt] = []) -> None: + ... + @property + def color_flows(self) -> list[list[list[tuple[int, int]]]]: + ... @color_flows.setter - def color_flows( - self, - arg0: collections.abc.Sequence[ - collections.abc.Sequence[ - collections.abc.Sequence[tuple[typing.SupportsInt, typing.SupportsInt]] - ] - ], - ) -> None: ... - @property - def diagram_color_indices(self) -> list[list[list[int]]]: ... + def color_flows(self, arg0: collections.abc.Sequence[collections.abc.Sequence[collections.abc.Sequence[tuple[typing.SupportsInt, typing.SupportsInt]]]]) -> None: + ... + @property + def diagram_color_indices(self) -> list[list[list[int]]]: + ... @diagram_color_indices.setter - def diagram_color_indices( - self, - arg0: collections.abc.Sequence[ - collections.abc.Sequence[collections.abc.Sequence[typing.SupportsInt]] - ], - ) -> None: ... - @property - def diagram_indices(self) -> list[list[int]]: ... + def diagram_color_indices(self, arg0: collections.abc.Sequence[collections.abc.Sequence[collections.abc.Sequence[typing.SupportsInt]]]) -> None: + ... + @property + def diagram_indices(self) -> list[list[int]]: + ... @diagram_indices.setter - def diagram_indices( - self, - arg0: collections.abc.Sequence[collections.abc.Sequence[typing.SupportsInt]], - ) -> None: ... + def diagram_indices(self, arg0: collections.abc.Sequence[collections.abc.Sequence[typing.SupportsInt]]) -> None: + ... @property - def helicities(self) -> list[list[float]]: ... + def helicities(self) -> list[list[float]]: + ... @helicities.setter - def helicities( - self, - arg0: collections.abc.Sequence[collections.abc.Sequence[typing.SupportsFloat]], - ) -> None: ... + def helicities(self, arg0: collections.abc.Sequence[collections.abc.Sequence[typing.SupportsFloat]]) -> None: + ... @property - def matrix_flavor_indices(self) -> list[int]: ... + def matrix_flavor_indices(self) -> list[int]: + ... @matrix_flavor_indices.setter - def matrix_flavor_indices( - self, arg0: collections.abc.Sequence[typing.SupportsInt] - ) -> None: ... + def matrix_flavor_indices(self, arg0: collections.abc.Sequence[typing.SupportsInt]) -> None: + ... @property - def pdg_color_types(self) -> dict[int, int]: ... + def pdg_color_types(self) -> dict[int, int]: + ... @pdg_color_types.setter - def pdg_color_types( - self, arg0: collections.abc.Mapping[typing.SupportsInt, typing.SupportsInt] - ) -> None: ... + def pdg_color_types(self, arg0: collections.abc.Mapping[typing.SupportsInt, typing.SupportsInt]) -> None: + ... @property - def pdg_ids(self) -> list[list[list[int]]]: ... + def pdg_ids(self) -> list[list[list[int]]]: + ... @pdg_ids.setter - def pdg_ids( - self, - arg0: collections.abc.Sequence[ - collections.abc.Sequence[collections.abc.Sequence[typing.SupportsInt]] - ], - ) -> None: ... - @property - def permutations(self) -> list[list[list[int]]]: ... + def pdg_ids(self, arg0: collections.abc.Sequence[collections.abc.Sequence[collections.abc.Sequence[typing.SupportsInt]]]) -> None: + ... + @property + def permutations(self) -> list[list[list[int]]]: + ... @permutations.setter - def permutations( - self, - arg0: collections.abc.Sequence[ - collections.abc.Sequence[collections.abc.Sequence[typing.SupportsInt]] - ], - ) -> None: ... - @property - def process_id(self) -> int: ... + def permutations(self, arg0: collections.abc.Sequence[collections.abc.Sequence[collections.abc.Sequence[typing.SupportsInt]]]) -> None: + ... + @property + def process_id(self) -> int: + ... @process_id.setter - def process_id(self, arg0: typing.SupportsInt) -> None: ... + def process_id(self, arg0: typing.SupportsInt) -> None: + ... @property - def topologies(self) -> list[Topology]: ... + def topologies(self) -> list[Topology]: + ... @topologies.setter - def topologies(self, arg0: collections.abc.Sequence[Topology]) -> None: ... - + def topologies(self, arg0: collections.abc.Sequence[Topology]) -> None: + ... class TPropagatorMapping(Mapping): - def __init__( - self, - integration_order: collections.abc.Sequence[typing.SupportsInt], - invariant_power: typing.SupportsFloat = 0.0, - ) -> None: ... - + def __init__(self, integration_order: collections.abc.Sequence[typing.SupportsInt], invariant_power: typing.SupportsFloat = 0.0) -> None: + ... + def random_dim(self) -> int: + ... class Tensor: @staticmethod - def numpy(tensor): ... + def numpy(tensor): + ... @staticmethod - def torch(tensor): ... - def __dlpack__( - self, - stream: typing.SupportsInt | None = None, - max_version: tuple[typing.SupportsInt, typing.SupportsInt] | None = None, - dl_device: typing.SupportsInt | None = None, - copy: bool | None = None, - ) -> typing.Any: ... - def __dlpack_device__(self) -> tuple[int, int]: ... - + def torch(tensor): + ... + def __dlpack__(self, stream: typing.SupportsInt | None = None, max_version: tuple[typing.SupportsInt, typing.SupportsInt] | None = None, dl_device: typing.SupportsInt | None = None, copy: bool | None = None) -> typing.Any: + ... + def __dlpack_device__(self) -> tuple[int, int]: + ... class ThreeBodyDecay(Mapping): - def __init__(self, com: bool) -> None: ... - + def __init__(self, com: bool) -> None: + ... class Topology: @staticmethod - def topologies(diagram: Diagram) -> list[Topology]: ... - def __init__(self, diagram: Diagram) -> None: ... - def propagator_momentum_terms( - self, arg0: bool - ) -> list[tuple[list[int], float, float]]: ... + def topologies(diagram: Diagram) -> list[Topology]: + ... + def __init__(self, diagram: Diagram) -> None: + ... + def propagator_momentum_terms(self, arg0: bool) -> list[tuple[list[int], float, float]]: + ... @property - def decay_integration_order(self) -> list[int]: ... + def decay_integration_order(self) -> list[int]: + ... @property - def decays(self) -> list[Decay]: ... + def decays(self) -> list[Decay]: + ... @property - def incoming_masses(self) -> list[float]: ... + def incoming_masses(self) -> list[float]: + ... @property - def outgoing_indices(self) -> list[int]: ... + def outgoing_indices(self) -> list[int]: + ... @property - def outgoing_masses(self) -> list[float]: ... + def outgoing_masses(self) -> list[float]: + ... @property - def t_integration_order(self) -> list[int]: ... + def t_integration_order(self) -> list[int]: + ... @property - def t_propagator_count(self) -> int: ... + def t_propagator_count(self) -> int: + ... @property - def t_propagator_masses(self) -> list[float]: ... + def t_propagator_masses(self) -> list[float]: + ... @property - def t_propagator_widths(self) -> list[float]: ... - + def t_propagator_widths(self) -> list[float]: + ... class TwoBodyDecay(Mapping): - def __init__(self, com: bool) -> None: ... - + def __init__(self, com: bool) -> None: + ... class TwoToThreeParticleScattering(Mapping): - def __init__( - self, - t_invariant_power: typing.SupportsFloat = 0.0, - t_mass: typing.SupportsFloat = 0.0, - t_width: typing.SupportsFloat = 0.0, - s_invariant_power: typing.SupportsFloat = 0.0, - s_mass: typing.SupportsFloat = 0.0, - s_width: typing.SupportsFloat = 0.0, - ) -> None: ... - + def __init__(self, t_invariant_power: typing.SupportsFloat = 0.0, t_mass: typing.SupportsFloat = 0.0, t_width: typing.SupportsFloat = 0.0, s_invariant_power: typing.SupportsFloat = 0.0, s_mass: typing.SupportsFloat = 0.0, s_width: typing.SupportsFloat = 0.0) -> None: + ... class TwoToTwoParticleScattering(Mapping): - def __init__( - self, - com: bool, - invariant_power: typing.SupportsFloat = 0.0, - mass: typing.SupportsFloat = 0.0, - width: typing.SupportsFloat = 0.0, - ) -> None: ... - + def __init__(self, com: bool, invariant_power: typing.SupportsFloat = 0.0, mass: typing.SupportsFloat = 0.0, width: typing.SupportsFloat = 0.0) -> None: + ... class Type: @typing.overload - def __init__( - self, - dtype: DataType, - batch_size: BatchSize, - shape: collections.abc.Sequence[typing.SupportsInt], - ) -> None: ... + def __init__(self, dtype: DataType, batch_size: BatchSize, shape: collections.abc.Sequence[typing.SupportsInt]) -> None: + ... @typing.overload - def __init__( - self, batch_size_list: collections.abc.Sequence[BatchSize] - ) -> None: ... - def __repr__(self) -> str: ... - def __str__(self) -> str: ... + def __init__(self, batch_size_list: collections.abc.Sequence[BatchSize]) -> None: + ... + def __repr__(self) -> str: + ... + def __str__(self) -> str: + ... @property - def batch_size(self) -> BatchSize: ... + def batch_size(self) -> BatchSize: + ... @property - def dtype(self) -> DataType: ... + def dtype(self) -> DataType: + ... @property - def shape(self) -> list[int]: ... - + def shape(self) -> list[int]: + ... class Unweighter(FunctionGenerator): - def __init__(self, types: collections.abc.Sequence[Type]) -> None: ... - + def __init__(self, types: NamedTypes) -> None: + ... class Value: @typing.overload - def __init__(self, value: typing.SupportsInt) -> None: ... + def __init__(self, value: typing.SupportsInt) -> None: + ... @typing.overload - def __init__(self, value: typing.SupportsFloat) -> None: ... - def __repr__(self) -> str: ... - def __str__(self) -> str: ... + def __init__(self, value: typing.SupportsFloat) -> None: + ... + def __repr__(self) -> str: + ... + def __str__(self) -> str: + ... @property - def literal_value( - self, - ) -> int | float | tuple[list[int], list[int] | list[float]] | None: ... + def literal_value(self) -> int | float | tuple[list[int], list[int] | list[float]] | None: + ... @property - def local_index(self) -> int: ... + def local_index(self) -> int: + ... @property - def type(self) -> Type: ... - + def type(self) -> Type: + ... class VegasGridOptimizer: - def __init__( - self, context: Context, grid_name: str, damping: typing.SupportsFloat - ) -> None: ... - def add_data(self, values: typing.Any, counts: typing.Any) -> None: ... - def optimize(self) -> None: ... - + def __init__(self, contexts: collections.abc.Sequence[Context], grid_name: str, damping: typing.SupportsFloat) -> None: + ... + def add_data(self, values: typing.Any, counts: typing.Any) -> None: + ... + def optimize(self) -> None: + ... class VegasHistogram(FunctionGenerator): - def __init__( - self, dimension: typing.SupportsInt, bin_count: typing.SupportsInt - ) -> None: ... - + def __init__(self, dimension: typing.SupportsInt, bin_count: typing.SupportsInt) -> None: + ... class VegasMapping(Mapping): - def __init__( - self, - dimension: typing.SupportsInt, - bin_count: typing.SupportsInt, - prefix: str = "", - ) -> None: ... - def grid_name(self) -> str: ... - def initialize_globals(self, context: Context) -> None: ... - -def batch_float_array(count: typing.SupportsInt) -> Type: ... -def batch_four_vec_array(count: typing.SupportsInt) -> Type: ... -def cpu_device() -> Device: ... -def cuda_device() -> Device: ... -def default_context() -> Context: ... -def default_cuda_context() -> Context: ... -def default_hip_context() -> Context: ... -def format_progress( - progress: typing.SupportsFloat, width: typing.SupportsInt -) -> str: ... -def format_si_prefix(value: typing.SupportsFloat) -> str: ... -def format_with_error( - value: typing.SupportsFloat, error: typing.SupportsFloat -) -> str: ... -def hip_device() -> Device: ... -def initialize_vegas_grid(context: Context, grid_name: str) -> None: ... -def multichannel_batch_size(count: typing.SupportsInt) -> Type: ... -def set_lib_path(lib_path: str) -> None: ... -def set_simd_vector_size(vector_size: typing.SupportsInt) -> None: ... -def set_thread_count(new_count: typing.SupportsInt) -> None: ... - + def __init__(self, dimension: typing.SupportsInt, bin_count: typing.SupportsInt, prefix: str = '') -> None: + ... + def grid_name(self) -> str: + ... + def initialize_globals(self, context: Context) -> None: + ... +class Verbosity: + """ + Members: + + silent + + log + + pretty + """ + __members__: typing.ClassVar[dict[str, Verbosity]] # value = {'silent': , 'log': , 'pretty': } + log: typing.ClassVar[Verbosity] # value = + pretty: typing.ClassVar[Verbosity] # value = + silent: typing.ClassVar[Verbosity] # value = + def __eq__(self, other: typing.Any) -> bool: + ... + def __getstate__(self) -> int: + ... + def __hash__(self) -> int: + ... + def __index__(self) -> int: + ... + @typing.overload + def __init__(self, value: typing.SupportsInt) -> None: + ... + @typing.overload + def __init__(self, name: str) -> None: + ... + def __int__(self) -> int: + ... + def __ne__(self, other: typing.Any) -> bool: + ... + def __repr__(self) -> str: + ... + def __setstate__(self, state: typing.SupportsInt) -> None: + ... + def __str__(self) -> str: + ... + @property + def name(self) -> str: + ... + @property + def value(self) -> int: + ... +def batch_float_array(count: typing.SupportsInt) -> Type: + ... +def batch_four_vec_array(count: typing.SupportsInt) -> Type: + ... +def cpu_device() -> Device: + ... +def cuda_device(index: typing.SupportsInt = 0) -> Device: + ... +def default_context() -> Context: + ... +def default_cuda_context(index: typing.SupportsInt = 0) -> Context: + ... +def default_hip_context(index: typing.SupportsInt = 0) -> Context: + ... +def format_progress(progress: typing.SupportsFloat, width: typing.SupportsInt) -> str: + ... +def format_si_prefix(value: typing.SupportsFloat) -> str: + ... +def format_with_error(value: typing.SupportsFloat, error: typing.SupportsFloat) -> str: + ... +def hip_device(index: typing.SupportsInt = 0) -> Device: + ... +def initialize_vegas_grid(context: Context, grid_name: str) -> None: + ... +def multichannel_batch_size(count: typing.SupportsInt) -> Type: + ... +def set_lib_path(lib_path: str) -> None: + ... +def set_simd_vector_size(vector_size: typing.SupportsInt) -> None: + ... batch_float: Type # value = float[batch_size] batch_four_vec: Type # value = float[batch_size, 4] batch_int: Type # value = int[batch_size] @@ -2408,8 +2416,8 @@ batch_size: BatchSize # value = batch_size batch_sizes: DataType # value = float: DataType # value = int: DataType # value = -log: EventGeneratorVerbosity # value = -pretty: EventGeneratorVerbosity # value = -silent: EventGeneratorVerbosity # value = +log: Verbosity # value = +pretty: Verbosity # value = +silent: Verbosity # value = single_float: Type # value = float[1] single_int: Type # value = int[1] diff --git a/madspace/pyproject.toml b/madspace/pyproject.toml index 0c544e012e..4047a259a3 100644 --- a/madspace/pyproject.toml +++ b/madspace/pyproject.toml @@ -5,10 +5,11 @@ build-backend = "scikit_build_core.build" [project] name = "madspace" version = "0.1.3" -description="Fast and flexible phase space mappings" +description="Fast and flexible phase-space mappings" readme = "README.md" authors = [ { name = "Theo Heimel", email = "theo.heimel@uclouvain.be" }, + { name = "Ramon Winterhalder", email = "ramon.winterhalder@unimi.it" }, ] requires-python = ">=3.11" classifiers = [ @@ -18,6 +19,14 @@ classifiers = [ ] dependencies = [] +[project.optional-dependencies] +test = [ + "pytest", + "numpy", + "scipy", + "torch", +] + [tool.scikit-build] wheel.expand-macos-universal-tags = true minimum-version = "build-system.requires" diff --git a/madspace/src/compgraphs/instruction_set_mixin.inc b/madspace/src/compgraphs/instruction_set_mixin.inc index 2ed08f15b5..cdaf72561a 100644 --- a/madspace/src/compgraphs/instruction_set_mixin.inc +++ b/madspace/src/compgraphs/instruction_set_mixin.inc @@ -76,87 +76,93 @@ InstructionOwner instructions[] { mi("two_to_two_particle_scattering_inverse", 58, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}), mi("two_to_three_particle_scattering", 59, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}}), mi("two_to_three_particle_scattering_inverse", 60, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_int, false, {}, false}, {DataType::dt_float, false, {}, false}}), - mi("three_body_decay_com", 61, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}}), - mi("three_body_decay_com_inverse", 62, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}), - mi("three_body_decay", 63, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {4}, false}}, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}}), - mi("three_body_decay_inverse", 64, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}}), - mi("t_inv_min_max", 65, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}), - mi("t_inv_value_and_min_max", 66, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}), - mi("s23_min_max", 67, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}), - mi("s23_value_and_min_max", 68, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}), - mi("invariants_from_momenta", 69, true, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_float, false, {"m", "n"}, false}}, {{DataType::dt_float, false, {"m"}, false}}), - mi("sde2_channel_weights", 70, true, {{DataType::dt_float, false, {"m"}, false}, {DataType::dt_float, false, {"c", "n"}, false}, {DataType::dt_float, false, {"c", "n"}, false}, {DataType::dt_int, false, {"c", "n"}, false}}, {{DataType::dt_float, false, {"c"}, false}}), - mi("subchannel_weights", 71, true, {{DataType::dt_float, false, {"m"}, false}, {DataType::dt_float, false, {"c", "n"}, false}, {DataType::dt_float, false, {"c", "n"}, false}, {DataType::dt_int, false, {"c", "n"}, false}, {DataType::dt_int, false, {"c", "n"}, false}, {DataType::dt_int, true, {"g"}, false}}, {{DataType::dt_float, false, {"c"}, false}}), - mi("apply_subchannel_weights", 72, true, {{DataType::dt_float, false, {"c"}, false}, {DataType::dt_float, false, {"n"}, false}, {DataType::dt_int, false, {"d"}, false}, {DataType::dt_int, false, {"d"}, false}}, {{DataType::dt_float, false, {"d"}, false}}), - mi("pt_eta_phi_x", 73, true, {{DataType::dt_float, false, {"n+2", 4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {"3n+2"}, false}}), - mi("mirror_momenta", 74, true, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_int, false, {}, false}}, {{DataType::dt_float, false, {"n", 4}, false}}), - mi("momenta_to_x1x2", 75, true, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}), - mi("uniform_invariant", 76, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}), - mi("uniform_invariant_inverse", 77, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}), - mi("breit_wigner_invariant", 78, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}), - mi("breit_wigner_invariant_inverse", 79, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}), - mi("stable_invariant", 80, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}), - mi("stable_invariant_inverse", 81, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}), - mi("stable_invariant_nu", 82, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}), - mi("stable_invariant_nu_inverse", 83, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}), - mi("fast_rambo_massless", 84, true, {{DataType::dt_float, false, {"3n-4"}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {4}, false}}, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_float, false, {}, false}}), - mi("fast_rambo_massless_inverse", 85, true, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {"3n-4"}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}}), - mi("fast_rambo_massless_com", 86, true, {{DataType::dt_float, false, {"3n-4"}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_float, false, {}, false}}), - mi("fast_rambo_massive", 87, true, {{DataType::dt_float, false, {"3n-4"}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {4}, false}}, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_float, false, {}, false}}), - mi("fast_rambo_massive_inverse", 88, true, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_float, false, {"3n-4"}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}}), - mi("fast_rambo_massive_com", 89, true, {{DataType::dt_float, false, {"3n-4"}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_float, false, {}, false}}), - mi("cut_unphysical", 90, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}}), - mi("cut_one", 91, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}}), - mi("cut_all", 92, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}}), - mi("cut_any", 93, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}}), - mi("scale_transverse_energy", 94, true, {{DataType::dt_float, false, {"n", 4}, false}}, {{DataType::dt_float, false, {}, false}}), - mi("scale_transverse_mass", 95, true, {{DataType::dt_float, false, {"n", 4}, false}}, {{DataType::dt_float, false, {}, false}}), - mi("scale_half_transverse_mass", 96, true, {{DataType::dt_float, false, {"n", 4}, false}}, {{DataType::dt_float, false, {}, false}}), - mi("scale_partonic_energy", 97, true, {{DataType::dt_float, false, {"n", 4}, false}}, {{DataType::dt_float, false, {}, false}}), - mi("chili_forward", 98, true, {{DataType::dt_float, false, {"3n-2"}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_float, false, {"n+2", 4}, false}, {DataType::dt_float, false, {}, false}}), - mi("chili_inverse", 99, true, {{DataType::dt_float, false, {"n+2", 4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_float, false, {"3n-2"}, false}, {DataType::dt_float, false, {}, false}}), - InstructionOwner(new MatrixElementInstruction(100, true)), - mi("collect_channel_weights", 101, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_int, true, {"n"}, false}, {DataType::dt_int, true, {"c"}, true}}, {{DataType::dt_float, false, {"c"}, false}}), - mi("interpolate_pdf", 102, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_int, false, {"n"}, false}, {DataType::dt_float, true, {"a"}, false}, {DataType::dt_float, true, {"b"}, false}, {DataType::dt_float, true, {16, "c", "d"}, false}}, {{DataType::dt_float, false, {"n"}, false}}), - mi("interpolate_alpha_s", 103, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, true, {"b+1"}, false}, {DataType::dt_float, true, {4, "b"}, false}}, {{DataType::dt_float, false, {}, false}}), - mi("matmul", 104, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, true, {"m", "n"}, false}, {DataType::dt_float, true, {"m"}, false}}, {{DataType::dt_float, false, {"m"}, false}}), - mi("relu", 105, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), - mi("leaky_relu", 106, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), - mi("elu", 107, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), - mi("gelu", 108, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), - mi("sigmoid", 109, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), - mi("softplus", 110, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), - InstructionOwner(new RqsReshapeInstruction(111, true)), - mi("rqs_find_bin", 112, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n", "b"}, false}, {DataType::dt_float, false, {"n", "b"}, false}, {DataType::dt_float, false, {"n", "b+1"}, false}}, {{DataType::dt_float, false, {"n", 6}, false}}), - mi("rqs_forward", 113, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n", 6}, false}}, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n"}, false}}), - mi("rqs_inverse", 114, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n", 6}, false}}, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n"}, false}}), - mi("softmax", 115, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), - mi("softmax_prior", 116, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_float, false, {"n"}, false}}), - mi("sample_discrete", 117, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_int, true, {}, false}}, {{DataType::dt_int, false, {}, false}, {DataType::dt_float, false, {}, false}}), - mi("sample_discrete_inverse", 118, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_int, true, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}), - mi("sample_discrete_probs", 119, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_int, false, {}, false}, {DataType::dt_float, false, {}, false}}), - mi("sample_discrete_probs_inverse", 120, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}), - mi("discrete_histogram", 121, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_int, true, {"n"}, true}}, {{DataType::dt_float, true, {"n"}, false}, {DataType::dt_int, true, {"n"}, false}}), - mi("permute_momenta", 122, true, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_int, true, {"m", "n"}, false}, {DataType::dt_int, false, {}, false}}, {{DataType::dt_float, false, {"n", 4}, false}}), - mi("gather", 123, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_float, false, {}, false}}), - mi("gather_int", 124, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_int, false, {"n"}, false}}, {{DataType::dt_int, false, {}, false}}), - mi("select_int", 125, true, {{DataType::dt_int, false, {"n"}, false}, {DataType::dt_int, false, {"m"}, false}}, {{DataType::dt_int, false, {"m"}, false}}), - mi("select", 126, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_int, false, {"m"}, false}}, {{DataType::dt_float, false, {"m"}, false}}), - mi("select_vector", 127, true, {{DataType::dt_float, false, {"n", "k"}, false}, {DataType::dt_int, false, {"m"}, false}}, {{DataType::dt_float, false, {"m", "k"}, false}}), - mi("argsort", 128, true, {{DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_int, false, {"n"}, false}}), - mi("one_hot", 129, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_int, true, {"n"}, true}}, {{DataType::dt_float, false, {"n"}, false}}), - mi("madnis_abs_weight", 130, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}}), - mi("madnis_softclip", 131, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}}), - mi("madnis_variance", 132, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}}), - mi("madnis_single_channel_variance", 133, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}}), - mi("madnis_multi_channel_variance", 134, true, {{DataType::dt_float, false, {"c"}, false}, {DataType::dt_float, false, {"c"}, false}}, {{DataType::dt_float, false, {}, false}}), - InstructionOwner(new NonzeroInstruction(135, true)), - InstructionOwner(new BatchGatherInstruction(136, true)), - InstructionOwner(new BatchScatterInstruction(137, true)), - InstructionOwner(new RandomInstruction(138, true)), - InstructionOwner(new UnweightInstruction(139, true)), - mi("vegas_forward", 140, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, true, {"n", "b"}, false}}, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n"}, false}}), - mi("vegas_inverse", 141, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, true, {"n", "b"}, false}}, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n"}, false}}), - mi("vegas_histogram", 142, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_int, true, {"b"}, true}}, {{DataType::dt_float, true, {"n", "b"}, false}, {DataType::dt_int, true, {"n", "b"}, false}}), - mi("histogram", 143, true, {{DataType::dt_float, false, {std::monostate{}}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_int, true, {"b"}, true}}, {{DataType::dt_float, true, {"b+2"}, false}, {DataType::dt_float, true, {"b+2"}, false}}), + mi("double_t_scattering", 61, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}}), + mi("double_t_scattering_inverse", 62, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}), + mi("three_body_decay_com", 63, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}}), + mi("three_body_decay_com_inverse", 64, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}), + mi("three_body_decay", 65, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {4}, false}}, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}}), + mi("three_body_decay_inverse", 66, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}}), + mi("t_inv_min_max", 67, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}), + mi("t_inv_value_and_min_max", 68, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}), + mi("t1_inv_min_max_doublet", 69, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}), + mi("t1_inv_value_and_min_max_doublet", 70, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}), + mi("t2_inv_min_max_doublet", 71, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}), + mi("t2_inv_value_and_min_max_doublet", 72, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}), + mi("s23_min_max", 73, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}), + mi("s23_value_and_min_max", 74, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}), + mi("invariants_from_momenta", 75, true, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_float, false, {"m", "n"}, false}}, {{DataType::dt_float, false, {"m"}, false}}), + mi("sde2_channel_weights", 76, true, {{DataType::dt_float, false, {"m"}, false}, {DataType::dt_float, false, {"c", "n"}, false}, {DataType::dt_float, false, {"c", "n"}, false}, {DataType::dt_int, false, {"c", "n"}, false}}, {{DataType::dt_float, false, {"c"}, false}}), + mi("subchannel_weights", 77, true, {{DataType::dt_float, false, {"m"}, false}, {DataType::dt_float, false, {"c", "n"}, false}, {DataType::dt_float, false, {"c", "n"}, false}, {DataType::dt_int, false, {"c", "n"}, false}, {DataType::dt_int, false, {"c", "n"}, false}, {DataType::dt_int, true, {"g"}, false}}, {{DataType::dt_float, false, {"c"}, false}}), + mi("apply_subchannel_weights", 78, true, {{DataType::dt_float, false, {"c"}, false}, {DataType::dt_float, false, {"n"}, false}, {DataType::dt_int, false, {"d"}, false}, {DataType::dt_int, false, {"d"}, false}}, {{DataType::dt_float, false, {"d"}, false}}), + mi("pt_eta_phi_x", 79, true, {{DataType::dt_float, false, {"n+2", 4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {"3n+2"}, false}}), + mi("mirror_momenta", 80, true, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_int, false, {}, false}}, {{DataType::dt_float, false, {"n", 4}, false}}), + mi("momenta_to_x1x2", 81, true, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}), + mi("uniform_invariant", 82, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}), + mi("uniform_invariant_inverse", 83, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}), + mi("breit_wigner_invariant", 84, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}), + mi("breit_wigner_invariant_inverse", 85, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}), + mi("stable_invariant", 86, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}), + mi("stable_invariant_inverse", 87, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}), + mi("stable_invariant_nu", 88, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}), + mi("stable_invariant_nu_inverse", 89, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}), + mi("fast_rambo_massless", 90, true, {{DataType::dt_float, false, {"3n-4"}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {4}, false}}, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_float, false, {}, false}}), + mi("fast_rambo_massless_inverse", 91, true, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {"3n-4"}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}}), + mi("fast_rambo_massless_com", 92, true, {{DataType::dt_float, false, {"3n-4"}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_float, false, {}, false}}), + mi("fast_rambo_massive", 93, true, {{DataType::dt_float, false, {"3n-4"}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {4}, false}}, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_float, false, {}, false}}), + mi("fast_rambo_massive_inverse", 94, true, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_float, false, {"3n-4"}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}}), + mi("fast_rambo_massive_com", 95, true, {{DataType::dt_float, false, {"3n-4"}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_float, false, {}, false}}), + mi("cut_unphysical", 96, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}}), + mi("cut_one", 97, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}}), + mi("cut_all", 98, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}}), + mi("cut_any", 99, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}}), + mi("scale_transverse_energy", 100, true, {{DataType::dt_float, false, {"n", 4}, false}}, {{DataType::dt_float, false, {}, false}}), + mi("scale_transverse_mass", 101, true, {{DataType::dt_float, false, {"n", 4}, false}}, {{DataType::dt_float, false, {}, false}}), + mi("scale_half_transverse_mass", 102, true, {{DataType::dt_float, false, {"n", 4}, false}}, {{DataType::dt_float, false, {}, false}}), + mi("scale_partonic_energy", 103, true, {{DataType::dt_float, false, {"n", 4}, false}}, {{DataType::dt_float, false, {}, false}}), + mi("chili_forward", 104, true, {{DataType::dt_float, false, {"3n-2"}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_float, false, {"n+2", 4}, false}, {DataType::dt_float, false, {}, false}}), + mi("chili_inverse", 105, true, {{DataType::dt_float, false, {"n+2", 4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_float, false, {"3n-2"}, false}, {DataType::dt_float, false, {}, false}}), + InstructionOwner(new MatrixElementInstruction(106, true)), + mi("collect_channel_weights", 107, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_int, true, {"n"}, false}, {DataType::dt_int, true, {"c"}, true}}, {{DataType::dt_float, false, {"c"}, false}}), + mi("interpolate_pdf", 108, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_int, false, {"n"}, false}, {DataType::dt_float, true, {"a"}, false}, {DataType::dt_float, true, {"b"}, false}, {DataType::dt_float, true, {16, "c", "d"}, false}}, {{DataType::dt_float, false, {"n"}, false}}), + mi("interpolate_alpha_s", 109, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, true, {"b+1"}, false}, {DataType::dt_float, true, {4, "b"}, false}}, {{DataType::dt_float, false, {}, false}}), + mi("matmul", 110, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, true, {"m", "n"}, false}, {DataType::dt_float, true, {"m"}, false}}, {{DataType::dt_float, false, {"m"}, false}}), + mi("relu", 111, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), + mi("leaky_relu", 112, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), + mi("elu", 113, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), + mi("gelu", 114, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), + mi("sigmoid", 115, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), + mi("softplus", 116, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), + InstructionOwner(new RqsReshapeInstruction(117, true)), + mi("rqs_find_bin", 118, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n", "b"}, false}, {DataType::dt_float, false, {"n", "b"}, false}, {DataType::dt_float, false, {"n", "b+1"}, false}}, {{DataType::dt_float, false, {"n", 6}, false}}), + mi("rqs_forward", 119, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n", 6}, false}}, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n"}, false}}), + mi("rqs_inverse", 120, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n", 6}, false}}, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n"}, false}}), + mi("softmax", 121, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), + mi("softmax_prior", 122, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_float, false, {"n"}, false}}), + mi("sample_discrete", 123, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_int, true, {}, false}}, {{DataType::dt_int, false, {}, false}, {DataType::dt_float, false, {}, false}}), + mi("sample_discrete_inverse", 124, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_int, true, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}), + mi("sample_discrete_probs", 125, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_int, false, {}, false}, {DataType::dt_float, false, {}, false}}), + mi("sample_discrete_probs_inverse", 126, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}), + mi("discrete_histogram", 127, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_int, true, {"n"}, true}}, {{DataType::dt_float, true, {"n"}, false}, {DataType::dt_int, true, {"n"}, false}}), + mi("permute_momenta", 128, true, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_int, true, {"m", "n"}, false}, {DataType::dt_int, false, {}, false}}, {{DataType::dt_float, false, {"n", 4}, false}}), + mi("gather", 129, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_float, false, {}, false}}), + mi("gather_int", 130, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_int, false, {"n"}, false}}, {{DataType::dt_int, false, {}, false}}), + mi("select_int", 131, true, {{DataType::dt_int, false, {"n"}, false}, {DataType::dt_int, false, {"m"}, false}}, {{DataType::dt_int, false, {"m"}, false}}), + mi("select", 132, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_int, false, {"m"}, false}}, {{DataType::dt_float, false, {"m"}, false}}), + mi("select_vector", 133, true, {{DataType::dt_float, false, {"n", "k"}, false}, {DataType::dt_int, false, {"m"}, false}}, {{DataType::dt_float, false, {"m", "k"}, false}}), + mi("argsort", 134, true, {{DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_int, false, {"n"}, false}}), + mi("one_hot", 135, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_int, true, {"n"}, true}}, {{DataType::dt_float, false, {"n"}, false}}), + mi("madnis_abs_weight", 136, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}}), + mi("madnis_softclip", 137, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}}), + mi("madnis_variance", 138, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}}), + mi("madnis_single_channel_variance", 139, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}}), + mi("madnis_multi_channel_variance", 140, true, {{DataType::dt_float, false, {"c"}, false}, {DataType::dt_float, false, {"c"}, false}}, {{DataType::dt_float, false, {}, false}}), + InstructionOwner(new NonzeroInstruction(141, true)), + InstructionOwner(new BatchGatherInstruction(142, true)), + InstructionOwner(new BatchScatterInstruction(143, true)), + InstructionOwner(new RandomInstruction(144, true)), + InstructionOwner(new UnweightInstruction(145, true)), + mi("vegas_forward", 146, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, true, {"n", "b"}, false}}, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n"}, false}}), + mi("vegas_inverse", 147, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, true, {"n", "b"}, false}}, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n"}, false}}), + mi("vegas_histogram", 148, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_int, true, {"b"}, true}}, {{DataType::dt_float, true, {"n", "b"}, false}, {DataType::dt_int, true, {"n", "b"}, false}}), + mi("histogram", 149, true, {{DataType::dt_float, false, {std::monostate{}}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_int, true, {"b"}, true}}, {{DataType::dt_float, true, {"b+2"}, false}, {DataType::dt_float, true, {"b+2"}, false}}), }; diff --git a/madspace/src/cpu/runtime_backward_mixin.inc b/madspace/src/cpu/runtime_backward_mixin.inc index 4474225a4d..9a1327b88b 100644 --- a/madspace/src/cpu/runtime_backward_mixin.inc +++ b/madspace/src/cpu/runtime_backward_mixin.inc @@ -52,66 +52,66 @@ case 24: case 25: backward_batch_foreach, backward_kernel_square, 2, 1, DeviceType>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 104: +case 110: backward_op_matmul(instr, locals, local_grads, device); break; -case 105: +case 111: backward_batch_foreach, backward_kernel_relu, 2, 1, DeviceType>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 106: +case 112: backward_batch_foreach, backward_kernel_leaky_relu, 2, 1, DeviceType>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 107: +case 113: backward_batch_foreach, backward_kernel_elu, 2, 1, DeviceType>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 108: +case 114: backward_batch_foreach, backward_kernel_gelu, 2, 1, DeviceType>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 109: +case 115: backward_batch_foreach, backward_kernel_sigmoid, 2, 1, DeviceType>, 1, 1, 0, 1>(instr, locals, local_grads, {}, {0}, device); break; -case 110: +case 116: backward_batch_foreach, backward_kernel_softplus, 2, 1, DeviceType>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 111: +case 117: backward_op_rqs_reshape(instr, locals, local_grads, device); break; -case 112: +case 118: backward_batch_foreach, backward_kernel_rqs_find_bin, 5, 4, 2, DeviceType>, 4, 1, 4, 0>(instr, locals, local_grads, {0,1,2,3}, {}, device); break; -case 113: +case 119: backward_batch_foreach, backward_kernel_rqs_forward, 4, 2, 2, DeviceType>, 2, 2, 2, 0>(instr, locals, local_grads, {0,1}, {}, device); break; -case 114: +case 120: backward_batch_foreach, backward_kernel_rqs_inverse, 4, 2, 2, DeviceType>, 2, 2, 2, 0>(instr, locals, local_grads, {0,1}, {}, device); break; -case 115: +case 121: backward_batch_foreach, backward_kernel_softmax, 2, 1, DeviceType>, 1, 1, 0, 1>(instr, locals, local_grads, {}, {0}, device); break; -case 116: +case 122: backward_batch_foreach, backward_kernel_softmax_prior, 2, 2, 1, DeviceType>, 2, 1, 0, 1>(instr, locals, local_grads, {}, {0}, device); break; -case 120: +case 126: backward_batch_foreach, backward_kernel_sample_discrete_probs_inverse, 4, 2, 1, DeviceType>, 2, 2, 2, 0>(instr, locals, local_grads, {0,1}, {}, device); break; -case 123: +case 129: backward_batch_foreach, backward_kernel_gather, 2, 2, 1, DeviceType>, 2, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 126: +case 132: backward_batch_foreach, backward_kernel_select, 2, 2, 1, DeviceType>, 2, 1, 1, 0>(instr, locals, local_grads, {1}, {}, device); break; -case 130: +case 136: backward_batch_foreach, backward_kernel_madnis_abs_weight, 3, 2, 1, DeviceType>, 2, 1, 2, 0>(instr, locals, local_grads, {0,1}, {}, device); break; -case 131: +case 137: backward_batch_foreach, backward_kernel_madnis_softclip, 5, 4, 1, DeviceType>, 4, 1, 4, 0>(instr, locals, local_grads, {0,1,2,3}, {}, device); break; -case 132: +case 138: backward_batch_foreach, backward_kernel_madnis_variance, 5, 4, 1, DeviceType>, 4, 1, 4, 0>(instr, locals, local_grads, {0,1,2,3}, {}, device); break; -case 133: +case 139: backward_batch_foreach, backward_kernel_madnis_single_channel_variance, 2, 2, 1, DeviceType>, 2, 1, 1, 0>(instr, locals, local_grads, {1}, {}, device); break; -case 134: +case 140: backward_batch_foreach, backward_kernel_madnis_multi_channel_variance, 3, 2, 1, DeviceType>, 2, 1, 2, 0>(instr, locals, local_grads, {0,1}, {}, device); break; diff --git a/madspace/src/cpu/runtime_mixin.inc b/madspace/src/cpu/runtime_mixin.inc index dea3f9934c..90fdeac07f 100644 --- a/madspace/src/cpu/runtime_mixin.inc +++ b/madspace/src/cpu/runtime_mixin.inc @@ -185,251 +185,269 @@ case 60: batch_foreach, kernel_two_to_three_particle_scattering_inverse, 7, 4, 1, DeviceType>, 7, 4>(instr, locals, device); break; case 61: - batch_foreach, kernel_three_body_decay_com, 9, 4, 1, DeviceType>, 9, 4>(instr, locals, device); + batch_foreach, kernel_double_t_scattering, 6, 3, 1, DeviceType>, 6, 3>(instr, locals, device); break; case 62: - batch_foreach, kernel_three_body_decay_com_inverse, 3, 10, 1, DeviceType>, 3, 10>(instr, locals, device); + batch_foreach, kernel_double_t_scattering_inverse, 4, 2, 1, DeviceType>, 4, 2>(instr, locals, device); break; case 63: - batch_foreach, kernel_three_body_decay, 10, 4, 1, DeviceType>, 10, 4>(instr, locals, device); + batch_foreach, kernel_three_body_decay_com, 9, 4, 1, DeviceType>, 9, 4>(instr, locals, device); break; case 64: - batch_foreach, kernel_three_body_decay_inverse, 3, 11, 1, DeviceType>, 3, 11>(instr, locals, device); + batch_foreach, kernel_three_body_decay_com_inverse, 3, 10, 1, DeviceType>, 3, 10>(instr, locals, device); break; case 65: - batch_foreach, kernel_t_inv_min_max, 4, 2, 1, DeviceType>, 4, 2>(instr, locals, device); + batch_foreach, kernel_three_body_decay, 10, 4, 1, DeviceType>, 10, 4>(instr, locals, device); break; case 66: - batch_foreach, kernel_t_inv_value_and_min_max, 4, 3, 1, DeviceType>, 4, 3>(instr, locals, device); + batch_foreach, kernel_three_body_decay_inverse, 3, 11, 1, DeviceType>, 3, 11>(instr, locals, device); break; case 67: - batch_foreach, kernel_s23_min_max, 6, 2, 1, DeviceType>, 6, 2>(instr, locals, device); + batch_foreach, kernel_t_inv_min_max, 4, 2, 1, DeviceType>, 4, 2>(instr, locals, device); break; case 68: - batch_foreach, kernel_s23_value_and_min_max, 6, 3, 1, DeviceType>, 6, 3>(instr, locals, device); + batch_foreach, kernel_t_inv_value_and_min_max, 4, 3, 1, DeviceType>, 4, 3>(instr, locals, device); break; case 69: - batch_foreach, kernel_invariants_from_momenta, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); + batch_foreach, kernel_t1_inv_min_max_doublet, 4, 2, 1, DeviceType>, 4, 2>(instr, locals, device); break; case 70: - batch_foreach, kernel_sde2_channel_weights, 4, 1, 1, DeviceType>, 4, 1>(instr, locals, device); + batch_foreach, kernel_t1_inv_value_and_min_max_doublet, 5, 3, 1, DeviceType>, 5, 3>(instr, locals, device); break; case 71: - batch_foreach, kernel_subchannel_weights, 6, 1, 1, DeviceType>, 6, 1>(instr, locals, device); + batch_foreach, kernel_t2_inv_min_max_doublet, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device); break; case 72: - batch_foreach, kernel_apply_subchannel_weights, 4, 1, 1, DeviceType>, 4, 1>(instr, locals, device); + batch_foreach, kernel_t2_inv_value_and_min_max_doublet, 6, 3, 1, DeviceType>, 6, 3>(instr, locals, device); break; case 73: - batch_foreach, kernel_pt_eta_phi_x, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device); + batch_foreach, kernel_s23_min_max, 6, 2, 1, DeviceType>, 6, 2>(instr, locals, device); break; case 74: - batch_foreach, kernel_mirror_momenta, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); + batch_foreach, kernel_s23_value_and_min_max, 6, 3, 1, DeviceType>, 6, 3>(instr, locals, device); break; case 75: - batch_foreach, kernel_momenta_to_x1x2, 2, 2, 1, DeviceType>, 2, 2>(instr, locals, device); + batch_foreach, kernel_invariants_from_momenta, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 76: - batch_foreach, kernel_uniform_invariant, 3, 2, 1, DeviceType>, 3, 2>(instr, locals, device); + batch_foreach, kernel_sde2_channel_weights, 4, 1, 1, DeviceType>, 4, 1>(instr, locals, device); break; case 77: - batch_foreach, kernel_uniform_invariant_inverse, 3, 2, 1, DeviceType>, 3, 2>(instr, locals, device); + batch_foreach, kernel_subchannel_weights, 6, 1, 1, DeviceType>, 6, 1>(instr, locals, device); break; case 78: - batch_foreach, kernel_breit_wigner_invariant, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device); + batch_foreach, kernel_apply_subchannel_weights, 4, 1, 1, DeviceType>, 4, 1>(instr, locals, device); break; case 79: - batch_foreach, kernel_breit_wigner_invariant_inverse, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device); + batch_foreach, kernel_pt_eta_phi_x, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device); break; case 80: - batch_foreach, kernel_stable_invariant, 4, 2, 1, DeviceType>, 4, 2>(instr, locals, device); + batch_foreach, kernel_mirror_momenta, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 81: - batch_foreach, kernel_stable_invariant_inverse, 4, 2, 1, DeviceType>, 4, 2>(instr, locals, device); + batch_foreach, kernel_momenta_to_x1x2, 2, 2, 1, DeviceType>, 2, 2>(instr, locals, device); break; case 82: - batch_foreach, kernel_stable_invariant_nu, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device); + batch_foreach, kernel_uniform_invariant, 3, 2, 1, DeviceType>, 3, 2>(instr, locals, device); break; case 83: - batch_foreach, kernel_stable_invariant_nu_inverse, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device); + batch_foreach, kernel_uniform_invariant_inverse, 3, 2, 1, DeviceType>, 3, 2>(instr, locals, device); break; case 84: - batch_foreach, kernel_fast_rambo_massless, 3, 2, 1, DeviceType>, 3, 2>(instr, locals, device); + batch_foreach, kernel_breit_wigner_invariant, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device); break; case 85: - batch_foreach, kernel_fast_rambo_massless_inverse, 2, 3, 1, DeviceType>, 2, 3>(instr, locals, device); + batch_foreach, kernel_breit_wigner_invariant_inverse, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device); break; case 86: - batch_foreach, kernel_fast_rambo_massless_com, 2, 2, 1, DeviceType>, 2, 2>(instr, locals, device); + batch_foreach, kernel_stable_invariant, 4, 2, 1, DeviceType>, 4, 2>(instr, locals, device); break; case 87: - batch_foreach, kernel_fast_rambo_massive, 4, 2, 1, DeviceType>, 4, 2>(instr, locals, device); + batch_foreach, kernel_stable_invariant_inverse, 4, 2, 1, DeviceType>, 4, 2>(instr, locals, device); break; case 88: - batch_foreach, kernel_fast_rambo_massive_inverse, 3, 3, 1, DeviceType>, 3, 3>(instr, locals, device); + batch_foreach, kernel_stable_invariant_nu, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device); break; case 89: - batch_foreach, kernel_fast_rambo_massive_com, 3, 2, 1, DeviceType>, 3, 2>(instr, locals, device); + batch_foreach, kernel_stable_invariant_nu_inverse, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device); break; case 90: - batch_foreach, kernel_cut_unphysical, 4, 1, 1, DeviceType>, 4, 1>(instr, locals, device); + batch_foreach, kernel_fast_rambo_massless, 3, 2, 1, DeviceType>, 3, 2>(instr, locals, device); break; case 91: - batch_foreach, kernel_cut_one, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device); + batch_foreach, kernel_fast_rambo_massless_inverse, 2, 3, 1, DeviceType>, 2, 3>(instr, locals, device); break; case 92: - batch_foreach, kernel_cut_all, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device); + batch_foreach, kernel_fast_rambo_massless_com, 2, 2, 1, DeviceType>, 2, 2>(instr, locals, device); break; case 93: - batch_foreach, kernel_cut_any, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device); + batch_foreach, kernel_fast_rambo_massive, 4, 2, 1, DeviceType>, 4, 2>(instr, locals, device); break; case 94: - batch_foreach, kernel_scale_transverse_energy, 1, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_fast_rambo_massive_inverse, 3, 3, 1, DeviceType>, 3, 3>(instr, locals, device); break; case 95: - batch_foreach, kernel_scale_transverse_mass, 1, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_fast_rambo_massive_com, 3, 2, 1, DeviceType>, 3, 2>(instr, locals, device); break; case 96: - batch_foreach, kernel_scale_half_transverse_mass, 1, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_cut_unphysical, 4, 1, 1, DeviceType>, 4, 1>(instr, locals, device); break; case 97: - batch_foreach, kernel_scale_partonic_energy, 1, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_cut_one, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device); break; case 98: - batch_foreach, kernel_chili_forward, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device); + batch_foreach, kernel_cut_all, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device); break; case 99: - batch_foreach, kernel_chili_inverse, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device); + batch_foreach, kernel_cut_any, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device); break; case 100: - op_matrix_element(instr, locals, device); + batch_foreach, kernel_scale_transverse_energy, 1, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 101: - batch_foreach, kernel_collect_channel_weights, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device); + batch_foreach, kernel_scale_transverse_mass, 1, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 102: - batch_foreach, kernel_interpolate_pdf, 6, 1, 1, DeviceType>, 6, 1>(instr, locals, device); + batch_foreach, kernel_scale_half_transverse_mass, 1, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 103: - batch_foreach, kernel_interpolate_alpha_s, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device); + batch_foreach, kernel_scale_partonic_energy, 1, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 104: - op_matmul(instr, locals, device); + batch_foreach, kernel_chili_forward, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device); break; case 105: - batch_foreach, kernel_relu, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_chili_inverse, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device); break; case 106: - batch_foreach, kernel_leaky_relu, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + op_matrix_element(instr, locals, device); break; case 107: - batch_foreach, kernel_elu, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_collect_channel_weights, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device); break; case 108: - batch_foreach, kernel_gelu, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_interpolate_pdf, 6, 1, 1, DeviceType>, 6, 1>(instr, locals, device); break; case 109: - batch_foreach, kernel_sigmoid, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_interpolate_alpha_s, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device); break; case 110: - batch_foreach, kernel_softplus, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + op_matmul(instr, locals, device); break; case 111: - op_rqs_reshape(instr, locals, device); + batch_foreach, kernel_relu, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 112: - batch_foreach, kernel_rqs_find_bin, 4, 1, 2, DeviceType>, 4, 1>(instr, locals, device); + batch_foreach, kernel_leaky_relu, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 113: - batch_foreach, kernel_rqs_forward, 2, 2, 2, DeviceType>, 2, 2>(instr, locals, device); + batch_foreach, kernel_elu, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 114: - batch_foreach, kernel_rqs_inverse, 2, 2, 2, DeviceType>, 2, 2>(instr, locals, device); + batch_foreach, kernel_gelu, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 115: - batch_foreach, kernel_softmax, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_sigmoid, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 116: - batch_foreach, kernel_softmax_prior, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); + batch_foreach, kernel_softplus, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 117: - batch_foreach, kernel_sample_discrete, 2, 2, 1, DeviceType>, 2, 2>(instr, locals, device); + op_rqs_reshape(instr, locals, device); break; case 118: - batch_foreach, kernel_sample_discrete_inverse, 2, 2, 1, DeviceType>, 2, 2>(instr, locals, device); + batch_foreach, kernel_rqs_find_bin, 4, 1, 2, DeviceType>, 4, 1>(instr, locals, device); break; case 119: - batch_foreach, kernel_sample_discrete_probs, 2, 2, 1, DeviceType>, 2, 2>(instr, locals, device); + batch_foreach, kernel_rqs_forward, 2, 2, 2, DeviceType>, 2, 2>(instr, locals, device); break; case 120: - batch_foreach, kernel_sample_discrete_probs_inverse, 2, 2, 1, DeviceType>, 2, 2>(instr, locals, device); + batch_foreach, kernel_rqs_inverse, 2, 2, 2, DeviceType>, 2, 2>(instr, locals, device); break; case 121: - op_discrete_histogram(instr, locals, device); + batch_foreach, kernel_softmax, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 122: - batch_foreach, kernel_permute_momenta, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device); + batch_foreach, kernel_softmax_prior, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 123: - batch_foreach, kernel_gather, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); + batch_foreach, kernel_sample_discrete, 2, 2, 1, DeviceType>, 2, 2>(instr, locals, device); break; case 124: - batch_foreach, kernel_gather_int, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); + batch_foreach, kernel_sample_discrete_inverse, 2, 2, 1, DeviceType>, 2, 2>(instr, locals, device); break; case 125: - batch_foreach, kernel_select_int, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); + batch_foreach, kernel_sample_discrete_probs, 2, 2, 1, DeviceType>, 2, 2>(instr, locals, device); break; case 126: - batch_foreach, kernel_select, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); + batch_foreach, kernel_sample_discrete_probs_inverse, 2, 2, 1, DeviceType>, 2, 2>(instr, locals, device); break; case 127: - batch_foreach, kernel_select_vector, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); + op_discrete_histogram(instr, locals, device); break; case 128: - batch_foreach, kernel_argsort, 1, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_permute_momenta, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device); break; case 129: - batch_foreach, kernel_one_hot, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); + batch_foreach, kernel_gather, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 130: - batch_foreach, kernel_madnis_abs_weight, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); + batch_foreach, kernel_gather_int, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 131: - batch_foreach, kernel_madnis_softclip, 4, 1, 1, DeviceType>, 4, 1>(instr, locals, device); + batch_foreach, kernel_select_int, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 132: - batch_foreach, kernel_madnis_variance, 4, 1, 1, DeviceType>, 4, 1>(instr, locals, device); + batch_foreach, kernel_select, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 133: - batch_foreach, kernel_madnis_single_channel_variance, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); + batch_foreach, kernel_select_vector, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 134: - batch_foreach, kernel_madnis_multi_channel_variance, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); + batch_foreach, kernel_argsort, 1, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 135: - op_nonzero(instr, locals, device); + batch_foreach, kernel_one_hot, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 136: - op_batch_gather(instr, locals, device); + batch_foreach, kernel_madnis_abs_weight, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 137: - op_batch_scatter(instr, locals, device); + batch_foreach, kernel_madnis_softclip, 4, 1, 1, DeviceType>, 4, 1>(instr, locals, device); break; case 138: - op_random(instr, locals, device); + batch_foreach, kernel_madnis_variance, 4, 1, 1, DeviceType>, 4, 1>(instr, locals, device); break; case 139: - op_unweight(instr, locals, device); + batch_foreach, kernel_madnis_single_channel_variance, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 140: - batch_foreach, kernel_vegas_forward, 2, 2, 2, DeviceType>, 2, 2>(instr, locals, device); + batch_foreach, kernel_madnis_multi_channel_variance, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 141: - batch_foreach, kernel_vegas_inverse, 2, 2, 2, DeviceType>, 2, 2>(instr, locals, device); + op_nonzero(instr, locals, device); break; case 142: - op_vegas_histogram(instr, locals, device); + op_batch_gather(instr, locals, device); break; case 143: + op_batch_scatter(instr, locals, device); + break; +case 144: + op_random(instr, locals, device); + break; +case 145: + op_unweight(instr, locals, device); + break; +case 146: + batch_foreach, kernel_vegas_forward, 2, 2, 2, DeviceType>, 2, 2>(instr, locals, device); + break; +case 147: + batch_foreach, kernel_vegas_inverse, 2, 2, 2, DeviceType>, 2, 2>(instr, locals, device); + break; +case 148: + op_vegas_histogram(instr, locals, device); + break; +case 149: op_histogram(instr, locals, device); break; diff --git a/madspace/src/gpu/runtime_backward_mixin.inc b/madspace/src/gpu/runtime_backward_mixin.inc index c0253b0f79..030d741310 100644 --- a/madspace/src/gpu/runtime_backward_mixin.inc +++ b/madspace/src/gpu/runtime_backward_mixin.inc @@ -52,66 +52,66 @@ case 24: case 25: backward_batch_foreach, 2, 1>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 104: +case 110: backward_op_matmul(instr, locals, local_grads, device); break; -case 105: +case 111: backward_batch_foreach, 2, 1>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 106: +case 112: backward_batch_foreach, 2, 1>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 107: +case 113: backward_batch_foreach, 2, 1>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 108: +case 114: backward_batch_foreach, 2, 1>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 109: +case 115: backward_batch_foreach, 2, 1>, 1, 1, 0, 1>(instr, locals, local_grads, {}, {0}, device); break; -case 110: +case 116: backward_batch_foreach, 2, 1>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 111: +case 117: backward_op_rqs_reshape(instr, locals, local_grads, device); break; -case 112: +case 118: backward_batch_foreach, 5, 4, 2>, 4, 1, 4, 0>(instr, locals, local_grads, {0,1,2,3}, {}, device); break; -case 113: +case 119: backward_batch_foreach, 4, 2, 2>, 2, 2, 2, 0>(instr, locals, local_grads, {0,1}, {}, device); break; -case 114: +case 120: backward_batch_foreach, 4, 2, 2>, 2, 2, 2, 0>(instr, locals, local_grads, {0,1}, {}, device); break; -case 115: +case 121: backward_batch_foreach, 2, 1>, 1, 1, 0, 1>(instr, locals, local_grads, {}, {0}, device); break; -case 116: +case 122: backward_batch_foreach, 2, 2, 1>, 2, 1, 0, 1>(instr, locals, local_grads, {}, {0}, device); break; -case 120: +case 126: backward_batch_foreach, 4, 2, 1>, 2, 2, 2, 0>(instr, locals, local_grads, {0,1}, {}, device); break; -case 123: +case 129: backward_batch_foreach, 2, 2, 1>, 2, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 126: +case 132: backward_batch_foreach, 2, 2, 1>, 2, 1, 1, 0>(instr, locals, local_grads, {1}, {}, device); break; -case 130: +case 136: backward_batch_foreach, 3, 2, 1>, 2, 1, 2, 0>(instr, locals, local_grads, {0,1}, {}, device); break; -case 131: +case 137: backward_batch_foreach, 5, 4, 1>, 4, 1, 4, 0>(instr, locals, local_grads, {0,1,2,3}, {}, device); break; -case 132: +case 138: backward_batch_foreach, 5, 4, 1>, 4, 1, 4, 0>(instr, locals, local_grads, {0,1,2,3}, {}, device); break; -case 133: +case 139: backward_batch_foreach, 2, 2, 1>, 2, 1, 1, 0>(instr, locals, local_grads, {1}, {}, device); break; -case 134: +case 140: backward_batch_foreach, 3, 2, 1>, 2, 1, 2, 0>(instr, locals, local_grads, {0,1}, {}, device); break; diff --git a/madspace/src/gpu/runtime_mixin.inc b/madspace/src/gpu/runtime_mixin.inc index 220f3bad33..3f3927e29a 100644 --- a/madspace/src/gpu/runtime_mixin.inc +++ b/madspace/src/gpu/runtime_mixin.inc @@ -185,251 +185,269 @@ case 60: batch_foreach, 7, 4, 1>, 7, 4>(instr, locals, device); break; case 61: - batch_foreach, 9, 4, 1>, 9, 4>(instr, locals, device); + batch_foreach, 6, 3, 1>, 6, 3>(instr, locals, device); break; case 62: - batch_foreach, 3, 10, 1>, 3, 10>(instr, locals, device); + batch_foreach, 4, 2, 1>, 4, 2>(instr, locals, device); break; case 63: - batch_foreach, 10, 4, 1>, 10, 4>(instr, locals, device); + batch_foreach, 9, 4, 1>, 9, 4>(instr, locals, device); break; case 64: - batch_foreach, 3, 11, 1>, 3, 11>(instr, locals, device); + batch_foreach, 3, 10, 1>, 3, 10>(instr, locals, device); break; case 65: - batch_foreach, 4, 2, 1>, 4, 2>(instr, locals, device); + batch_foreach, 10, 4, 1>, 10, 4>(instr, locals, device); break; case 66: - batch_foreach, 4, 3, 1>, 4, 3>(instr, locals, device); + batch_foreach, 3, 11, 1>, 3, 11>(instr, locals, device); break; case 67: - batch_foreach, 6, 2, 1>, 6, 2>(instr, locals, device); + batch_foreach, 4, 2, 1>, 4, 2>(instr, locals, device); break; case 68: - batch_foreach, 6, 3, 1>, 6, 3>(instr, locals, device); + batch_foreach, 4, 3, 1>, 4, 3>(instr, locals, device); break; case 69: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + batch_foreach, 4, 2, 1>, 4, 2>(instr, locals, device); break; case 70: - batch_foreach, 4, 1, 1>, 4, 1>(instr, locals, device); + batch_foreach, 5, 3, 1>, 5, 3>(instr, locals, device); break; case 71: - batch_foreach, 6, 1, 1>, 6, 1>(instr, locals, device); + batch_foreach, 5, 2, 1>, 5, 2>(instr, locals, device); break; case 72: - batch_foreach, 4, 1, 1>, 4, 1>(instr, locals, device); + batch_foreach, 6, 3, 1>, 6, 3>(instr, locals, device); break; case 73: - batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); + batch_foreach, 6, 2, 1>, 6, 2>(instr, locals, device); break; case 74: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + batch_foreach, 6, 3, 1>, 6, 3>(instr, locals, device); break; case 75: - batch_foreach, 2, 2, 1>, 2, 2>(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 76: - batch_foreach, 3, 2, 1>, 3, 2>(instr, locals, device); + batch_foreach, 4, 1, 1>, 4, 1>(instr, locals, device); break; case 77: - batch_foreach, 3, 2, 1>, 3, 2>(instr, locals, device); + batch_foreach, 6, 1, 1>, 6, 1>(instr, locals, device); break; case 78: - batch_foreach, 5, 2, 1>, 5, 2>(instr, locals, device); + batch_foreach, 4, 1, 1>, 4, 1>(instr, locals, device); break; case 79: - batch_foreach, 5, 2, 1>, 5, 2>(instr, locals, device); + batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); break; case 80: - batch_foreach, 4, 2, 1>, 4, 2>(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 81: - batch_foreach, 4, 2, 1>, 4, 2>(instr, locals, device); + batch_foreach, 2, 2, 1>, 2, 2>(instr, locals, device); break; case 82: - batch_foreach, 5, 2, 1>, 5, 2>(instr, locals, device); + batch_foreach, 3, 2, 1>, 3, 2>(instr, locals, device); break; case 83: - batch_foreach, 5, 2, 1>, 5, 2>(instr, locals, device); + batch_foreach, 3, 2, 1>, 3, 2>(instr, locals, device); break; case 84: - batch_foreach, 3, 2, 1>, 3, 2>(instr, locals, device); + batch_foreach, 5, 2, 1>, 5, 2>(instr, locals, device); break; case 85: - batch_foreach, 2, 3, 1>, 2, 3>(instr, locals, device); + batch_foreach, 5, 2, 1>, 5, 2>(instr, locals, device); break; case 86: - batch_foreach, 2, 2, 1>, 2, 2>(instr, locals, device); + batch_foreach, 4, 2, 1>, 4, 2>(instr, locals, device); break; case 87: - batch_foreach, 4, 2, 1>, 4, 2>(instr, locals, device); + batch_foreach, 4, 2, 1>, 4, 2>(instr, locals, device); break; case 88: - batch_foreach, 3, 3, 1>, 3, 3>(instr, locals, device); + batch_foreach, 5, 2, 1>, 5, 2>(instr, locals, device); break; case 89: - batch_foreach, 3, 2, 1>, 3, 2>(instr, locals, device); + batch_foreach, 5, 2, 1>, 5, 2>(instr, locals, device); break; case 90: - batch_foreach, 4, 1, 1>, 4, 1>(instr, locals, device); + batch_foreach, 3, 2, 1>, 3, 2>(instr, locals, device); break; case 91: - batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); + batch_foreach, 2, 3, 1>, 2, 3>(instr, locals, device); break; case 92: - batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); + batch_foreach, 2, 2, 1>, 2, 2>(instr, locals, device); break; case 93: - batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); + batch_foreach, 4, 2, 1>, 4, 2>(instr, locals, device); break; case 94: - batch_foreach, 1, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 3, 3, 1>, 3, 3>(instr, locals, device); break; case 95: - batch_foreach, 1, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 3, 2, 1>, 3, 2>(instr, locals, device); break; case 96: - batch_foreach, 1, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 4, 1, 1>, 4, 1>(instr, locals, device); break; case 97: - batch_foreach, 1, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); break; case 98: - batch_foreach, 5, 2, 1>, 5, 2>(instr, locals, device); + batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); break; case 99: - batch_foreach, 5, 2, 1>, 5, 2>(instr, locals, device); + batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); break; case 100: - op_matrix_element(instr, locals, device); + batch_foreach, 1, 1, 1>, 1, 1>(instr, locals, device); break; case 101: - batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); + batch_foreach, 1, 1, 1>, 1, 1>(instr, locals, device); break; case 102: - batch_foreach, 6, 1, 1>, 6, 1>(instr, locals, device); + batch_foreach, 1, 1, 1>, 1, 1>(instr, locals, device); break; case 103: - batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); + batch_foreach, 1, 1, 1>, 1, 1>(instr, locals, device); break; case 104: - op_matmul(instr, locals, device); + batch_foreach, 5, 2, 1>, 5, 2>(instr, locals, device); break; case 105: - batch_foreach, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 5, 2, 1>, 5, 2>(instr, locals, device); break; case 106: - batch_foreach, 1, 1>, 1, 1>(instr, locals, device); + op_matrix_element(instr, locals, device); break; case 107: - batch_foreach, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); break; case 108: - batch_foreach, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 6, 1, 1>, 6, 1>(instr, locals, device); break; case 109: - batch_foreach, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); break; case 110: - batch_foreach, 1, 1>, 1, 1>(instr, locals, device); + op_matmul(instr, locals, device); break; case 111: - op_rqs_reshape(instr, locals, device); + batch_foreach, 1, 1>, 1, 1>(instr, locals, device); break; case 112: - batch_foreach, 4, 1, 2>, 4, 1>(instr, locals, device); + batch_foreach, 1, 1>, 1, 1>(instr, locals, device); break; case 113: - batch_foreach, 2, 2, 2>, 2, 2>(instr, locals, device); + batch_foreach, 1, 1>, 1, 1>(instr, locals, device); break; case 114: - batch_foreach, 2, 2, 2>, 2, 2>(instr, locals, device); + batch_foreach, 1, 1>, 1, 1>(instr, locals, device); break; case 115: - batch_foreach, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 1, 1>, 1, 1>(instr, locals, device); break; case 116: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + batch_foreach, 1, 1>, 1, 1>(instr, locals, device); break; case 117: - batch_foreach, 2, 2, 1>, 2, 2>(instr, locals, device); + op_rqs_reshape(instr, locals, device); break; case 118: - batch_foreach, 2, 2, 1>, 2, 2>(instr, locals, device); + batch_foreach, 4, 1, 2>, 4, 1>(instr, locals, device); break; case 119: - batch_foreach, 2, 2, 1>, 2, 2>(instr, locals, device); + batch_foreach, 2, 2, 2>, 2, 2>(instr, locals, device); break; case 120: - batch_foreach, 2, 2, 1>, 2, 2>(instr, locals, device); + batch_foreach, 2, 2, 2>, 2, 2>(instr, locals, device); break; case 121: - op_discrete_histogram(instr, locals, device); + batch_foreach, 1, 1>, 1, 1>(instr, locals, device); break; case 122: - batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 123: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + batch_foreach, 2, 2, 1>, 2, 2>(instr, locals, device); break; case 124: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + batch_foreach, 2, 2, 1>, 2, 2>(instr, locals, device); break; case 125: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + batch_foreach, 2, 2, 1>, 2, 2>(instr, locals, device); break; case 126: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + batch_foreach, 2, 2, 1>, 2, 2>(instr, locals, device); break; case 127: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + op_discrete_histogram(instr, locals, device); break; case 128: - batch_foreach, 1, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); break; case 129: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 130: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 131: - batch_foreach, 4, 1, 1>, 4, 1>(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 132: - batch_foreach, 4, 1, 1>, 4, 1>(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 133: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 134: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + batch_foreach, 1, 1, 1>, 1, 1>(instr, locals, device); break; case 135: - op_nonzero(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 136: - op_batch_gather(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 137: - op_batch_scatter(instr, locals, device); + batch_foreach, 4, 1, 1>, 4, 1>(instr, locals, device); break; case 138: - op_random(instr, locals, device); + batch_foreach, 4, 1, 1>, 4, 1>(instr, locals, device); break; case 139: - op_unweight(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 140: - batch_foreach, 2, 2, 2>, 2, 2>(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 141: - batch_foreach, 2, 2, 2>, 2, 2>(instr, locals, device); + op_nonzero(instr, locals, device); break; case 142: - op_vegas_histogram(instr, locals, device); + op_batch_gather(instr, locals, device); break; case 143: + op_batch_scatter(instr, locals, device); + break; +case 144: + op_random(instr, locals, device); + break; +case 145: + op_unweight(instr, locals, device); + break; +case 146: + batch_foreach, 2, 2, 2>, 2, 2>(instr, locals, device); + break; +case 147: + batch_foreach, 2, 2, 2>, 2, 2>(instr, locals, device); + break; +case 148: + op_vegas_histogram(instr, locals, device); + break; +case 149: op_histogram(instr, locals, device); break; diff --git a/madspace/src/kernels/kernels.hpp b/madspace/src/kernels/kernels.hpp index c436367728..7c2d828efc 100644 --- a/madspace/src/kernels/kernels.hpp +++ b/madspace/src/kernels/kernels.hpp @@ -5,6 +5,7 @@ #include "discrete.hpp" #include "invariants.hpp" #include "kinematics.hpp" +#include "lup_det.hpp" #include "math.hpp" #include "multichannel.hpp" #include "nn.hpp" diff --git a/madspace/src/kernels/kinematics.hpp b/madspace/src/kernels/kinematics.hpp index 9150233e46..a0ba7916f6 100644 --- a/madspace/src/kernels/kinematics.hpp +++ b/madspace/src/kernels/kinematics.hpp @@ -106,6 +106,29 @@ t_inv_min_max(FVal s, FVal ma_2, FVal mb_2, FVal m1_2, FVal m2_2) return {t_min, t_max}; } +template +KERNELSPEC Pair, FVal> +t1_inv_min_max_doublet(FVal s, FVal m1_2, FVal mir_min_2) { + // |t1| bounds for pa+pb -> p_i + p_ir with massless pa,pb, fixed m_i, + // and m_ir >= mir_min. Same shape as standard t bounds with + // ma=mb=0, m2=mir_min. + return t_inv_min_max(s, 0., 0., m1_2, mir_min_2); +} + +template +KERNELSPEC Pair, FVal> +t2_inv_min_max_doublet(FVal s, FVal m1_2, FVal mir_min_2, FVal t1_abs) { + // |t2| bounds given fixed s, m_i, mir_min, |t1|. + // Derivation (abs-value convention): from the signed-t formulas in the + // Fortran double_t, flip both signs. + auto denom = m1_2 + t1_abs + EPS; + auto t2_min_raw = m1_2 * (s - t1_abs - m1_2) / denom; + auto t2_max_raw = s - t1_abs - m1_2 - mir_min_2; + auto t2_min = max(t2_min_raw, 0.); + auto t2_max = where(t2_max_raw > t2_min, t2_max_raw, t2_min + EPS); + return {t2_min, t2_max}; +} + template KERNELSPEC FVal lsquare(FourMom p) { return p[0] * p[0] - p[1] * p[1] - p[2] * p[2] - p[3] * p[3]; @@ -197,6 +220,77 @@ KERNELSPEC FourMom rotate_inverse(FourMom p, FourMom q) { }; } +template +KERNELSPEC FourMom +rotate_two_ref(FourMom p, FourMom q_z, FourMom q_x) { + // Forward rotation: take p from a canonical frame + // (e_z along q_z, e_x in the plane spanned by q_z and q_x with positive + // component along q_x's perpendicular part) + // into the lab frame. Energy unchanged. + // + // Used by kernel_two_to_three_particle_scattering: q_z = pa_com, + // q_x = p3 boosted into the p_12 rest frame. Picks a unique azimuth + // so that the (q_z, q_x) plane is the phi=0 half-plane. + + // z_hat = q_z / |q_z| + auto qz_n2 = q_z[1]*q_z[1] + q_z[2]*q_z[2] + q_z[3]*q_z[3]; + auto qz_n = sqrt(max(qz_n2, EPS2)); + auto zx = q_z[1] / qz_n, zy = q_z[2] / qz_n, zz = q_z[3] / qz_n; + + // x_hat = (q_x perp to z_hat) / |...| + auto qx_dot_z = q_x[1]*zx + q_x[2]*zy + q_x[3]*zz; + auto rx = q_x[1] - qx_dot_z*zx; + auto ry = q_x[2] - qx_dot_z*zy; + auto rz = q_x[3] - qx_dot_z*zz; + auto rn2 = rx*rx + ry*ry + rz*rz; + auto rn = sqrt(max(rn2, EPS2)); + auto xx = rx / rn, xy = ry / rn, xz = rz / rn; + + // y_hat = z_hat x x_hat + auto yx = zy*xz - zz*xy; + auto yy = zz*xx - zx*xz; + auto yz = zx*xy - zy*xx; + + // world spatial = p[1]*x_hat + p[2]*y_hat + p[3]*z_hat + return FourMom{ + p[0], + p[1]*xx + p[2]*yx + p[3]*zx, + p[1]*xy + p[2]*yy + p[3]*zy, + p[1]*xz + p[2]*yz + p[3]*zz, + }; +} + +template +KERNELSPEC FourMom +rotate_two_ref_inverse(FourMom p, FourMom q_z, FourMom q_x) { + // Inverse of rotate_two_ref: take a vector p from the lab frame into the + // canonical frame defined by (q_z, q_x). Energy unchanged. + + auto qz_n2 = q_z[1]*q_z[1] + q_z[2]*q_z[2] + q_z[3]*q_z[3]; + auto qz_n = sqrt(max(qz_n2, EPS2)); + auto zx = q_z[1] / qz_n, zy = q_z[2] / qz_n, zz = q_z[3] / qz_n; + + auto qx_dot_z = q_x[1]*zx + q_x[2]*zy + q_x[3]*zz; + auto rx = q_x[1] - qx_dot_z*zx; + auto ry = q_x[2] - qx_dot_z*zy; + auto rz = q_x[3] - qx_dot_z*zz; + auto rn2 = rx*rx + ry*ry + rz*rz; + auto rn = sqrt(max(rn2, EPS2)); + auto xx = rx / rn, xy = ry / rn, xz = rz / rn; + + auto yx = zy*xz - zz*xy; + auto yy = zz*xx - zx*xz; + auto yz = zx*xy - zy*xx; + + // canonical components = world spatial . (x_hat, y_hat, z_hat) + return FourMom{ + p[0], + p[1]*xx + p[2]*xy + p[3]*xz, + p[1]*yx + p[2]*yy + p[3]*yz, + p[1]*zx + p[2]*zy + p[3]*zz, + }; +} + template KERNELSPEC FourMom boost(FourMom k, FourMom p_boost, FVal sign) { // Perform the boost @@ -413,6 +507,93 @@ KERNELSPEC void kernel_invariants_from_momenta( } } +template +KERNELSPEC void kernel_t1_inv_min_max_doublet( + FIn pa, + FIn pb, + FIn m1, + FIn mir_min, + FOut t_min, + FOut t_max +) { + FourMom p_tot; + for (int i = 0; i < 4; ++i) { + p_tot[i] = pa[i] + pb[i]; + } + auto s = lsquare(p_tot); + auto bounds = t1_inv_min_max_doublet(s, m1 * m1, mir_min * mir_min); + t_min = bounds.first; + t_max = bounds.second; +} + +template +KERNELSPEC void kernel_t1_inv_value_and_min_max_doublet( + FIn pa, + FIn pb, + FIn p1, + FIn m1, + FIn mir_min, + FOut t_abs, + FOut t_min, + FOut t_max +) { + FourMom pa1, p_tot; + for (int i = 0; i < 4; ++i) { + pa1[i] = pa[i] - p1[i]; + p_tot[i] = pa[i] + pb[i]; + } + auto s = lsquare(p_tot); + auto bounds = t1_inv_min_max_doublet(s, m1 * m1, mir_min * mir_min); + t_min = bounds.first; + t_max = bounds.second; + t_abs = -lsquare(pa1); +} + +template +KERNELSPEC void kernel_t2_inv_min_max_doublet( + FIn pa, + FIn pb, + FIn m1, + FIn mir_min, + FIn t1_abs, + FOut t_min, + FOut t_max +) { + FourMom p_tot; + for (int i = 0; i < 4; ++i) { + p_tot[i] = pa[i] + pb[i]; + } + auto s = lsquare(p_tot); + auto bounds = + t2_inv_min_max_doublet(s, m1 * m1, mir_min * mir_min, t1_abs); + t_min = bounds.first; + t_max = bounds.second; +} + +template +KERNELSPEC void kernel_t2_inv_value_and_min_max_doublet( + FIn pa, + FIn pb, + FIn p1, + FIn m1, + FIn mir_min, + FIn t1_abs, + FOut t_abs, + FOut t_min, + FOut t_max +) { + FourMom pb1, p_tot; + for (int i = 0; i < 4; ++i) { + pb1[i] = pb[i] - p1[i]; + p_tot[i] = pa[i] + pb[i]; + } + auto s = lsquare(p_tot); + auto bounds = t2_inv_min_max_doublet(s, m1 * m1, mir_min * mir_min, t1_abs); + t_min = bounds.first; + t_max = bounds.second; + t_abs = -lsquare(pb1); +} + template KERNELSPEC void kernel_pt_eta_phi_x(FIn p_ext, FIn x1, FIn x2, FOut output) { diff --git a/madspace/src/kernels/lup_det.hpp b/madspace/src/kernels/lup_det.hpp new file mode 100644 index 0000000000..da0b8b704b --- /dev/null +++ b/madspace/src/kernels/lup_det.hpp @@ -0,0 +1,371 @@ +#pragma once + +#include "definitions.hpp" + +namespace madspace { +namespace kernels { + +// Local tolerance constant, intentionally tighter than the module-level EPS. +// Chosen so the relative tolerance scale * LU_TOL sits above eps_mach but +// below the noise floor at which the polynomial expansion blows up. +inline constexpr double LU_TOL = 1e-14; + +// Variadic max-of-|.| over an arbitrary list of FVal. +template +KERNELSPEC FVal max_abs(FVal first, Args... rest) { + if constexpr (sizeof...(rest) == 0) { + return fabs(first); + } else { + return max(fabs(first), max_abs(rest...)); + } +} + +// det of a symmetric 3x3 matrix via LU with partial pivoting. +// +// [ a11 a12 a13 ] +// [ a12 a22 a23 ] +// [ a13 a23 a33 ] +// +// Falls back lane-wise to `polynomial_fallback` if any pivot is below +// LU_TOL * max(|entries|). +template +KERNELSPEC FVal lup_det3( + FVal a11, FVal a12, FVal a13, + FVal a22, FVal a23, + FVal a33, + FVal polynomial_fallback +) { + // Row-major working copy. + auto r00 = a11, r01 = a12, r02 = a13; + auto r10 = a12, r11 = a22, r12 = a23; + auto r20 = a13, r21 = a23, r22 = a33; + + auto scale = max_abs(a11, a12, a13, a22, a23, a33); + auto tol = LU_TOL * scale; + auto sign = FVal(1.0); + auto degenerate = (scale == 0.); + + // ---- Step 0: pivot from column 0, rows {0,1,2} ---- + { + auto m0 = fabs(r00), m1 = fabs(r10), m2 = fabs(r20); + auto best = IVal(0); + auto best_mag = m0; + auto pick1 = (m1 > best_mag); + best = where(pick1, IVal(1), best); + best_mag = where(pick1, m1, best_mag); + auto pick2 = (m2 > best_mag); + best = where(pick2, IVal(2), best); + best_mag = where(pick2, m2, best_mag); + + auto sel1 = (best == IVal(1)); + auto sel2 = (best == IVal(2)); + auto any_swap = sel1 | sel2; + + auto n00 = where(sel1, r10, where(sel2, r20, r00)); + auto n01 = where(sel1, r11, where(sel2, r21, r01)); + auto n02 = where(sel1, r12, where(sel2, r22, r02)); + auto n10 = where(sel1, r00, r10); + auto n11 = where(sel1, r01, r11); + auto n12 = where(sel1, r02, r12); + auto n20 = where(sel2, r00, r20); + auto n21 = where(sel2, r01, r21); + auto n22 = where(sel2, r02, r22); + r00 = n00; r01 = n01; r02 = n02; + r10 = n10; r11 = n11; r12 = n12; + r20 = n20; r21 = n21; r22 = n22; + + sign = where(any_swap, -sign, sign); + degenerate = degenerate | (best_mag < tol); + + auto pivot_safe = where(degenerate, FVal(1.0), r00); + auto f10 = r10 / pivot_safe; + auto f20 = r20 / pivot_safe; + r10 = f10; + r11 = r11 - f10 * r01; + r12 = r12 - f10 * r02; + r20 = f20; + r21 = r21 - f20 * r01; + r22 = r22 - f20 * r02; + } + + // ---- Step 1: pivot from column 1, rows {1, 2} ---- + { + auto m1 = fabs(r11), m2 = fabs(r21); + auto sel2 = (m2 > m1); + auto pivot_mag = where(sel2, m2, m1); + + auto n11 = where(sel2, r21, r11); + auto n12 = where(sel2, r22, r12); + auto n21 = where(sel2, r11, r21); + auto n22 = where(sel2, r12, r22); + r11 = n11; r12 = n12; r21 = n21; r22 = n22; + // Multipliers in column 0 follow the row swap (preserves L, not + // strictly needed for the determinant but kept for correctness). + auto n10 = where(sel2, r20, r10); + auto n20 = where(sel2, r10, r20); + r10 = n10; r20 = n20; + + sign = where(sel2, -sign, sign); + degenerate = degenerate | (pivot_mag < tol); + + auto pivot_safe = where(degenerate, FVal(1.0), r11); + auto f21 = r21 / pivot_safe; + r21 = f21; + r22 = r22 - f21 * r12; + } + + // ---- Final pivot check: the trailing diagonal entry is also a pivot ---- + degenerate = degenerate | (fabs(r22) < tol); + + auto lu_det = sign * r00 * r11 * r22; + return where(degenerate, polynomial_fallback, lu_det); +} + +// det of a general (non-symmetric) 3x3 matrix via LU with partial pivoting. +// +// [ a11 a12 a13 ] +// [ a21 a22 a23 ] +// [ a31 a32 a33 ] +// +// Same fallback policy as lup_det3. Needed for bk_V, whose matrix has +// a21 = a12 but a13 != a31 and a23 != a32 (and a22 = a33 = 0). +template +KERNELSPEC FVal lup_det3_general( + FVal a11, FVal a12, FVal a13, + FVal a21, FVal a22, FVal a23, + FVal a31, FVal a32, FVal a33, + FVal polynomial_fallback +) { + auto r00 = a11, r01 = a12, r02 = a13; + auto r10 = a21, r11 = a22, r12 = a23; + auto r20 = a31, r21 = a32, r22 = a33; + + auto scale = max_abs(a11, a12, a13, a21, a22, a23, a31, a32, a33); + auto tol = LU_TOL * scale; + auto sign = FVal(1.0); + auto degenerate = (scale == 0.); + + // ---- Step 0: pivot from column 0, rows {0,1,2} ---- + { + auto m0 = fabs(r00), m1 = fabs(r10), m2 = fabs(r20); + auto best = IVal(0); + auto best_mag = m0; + auto pick1 = (m1 > best_mag); + best = where(pick1, IVal(1), best); + best_mag = where(pick1, m1, best_mag); + auto pick2 = (m2 > best_mag); + best = where(pick2, IVal(2), best); + best_mag = where(pick2, m2, best_mag); + + auto sel1 = (best == IVal(1)); + auto sel2 = (best == IVal(2)); + auto any_swap = sel1 | sel2; + + auto n00 = where(sel1, r10, where(sel2, r20, r00)); + auto n01 = where(sel1, r11, where(sel2, r21, r01)); + auto n02 = where(sel1, r12, where(sel2, r22, r02)); + auto n10 = where(sel1, r00, r10); + auto n11 = where(sel1, r01, r11); + auto n12 = where(sel1, r02, r12); + auto n20 = where(sel2, r00, r20); + auto n21 = where(sel2, r01, r21); + auto n22 = where(sel2, r02, r22); + r00 = n00; r01 = n01; r02 = n02; + r10 = n10; r11 = n11; r12 = n12; + r20 = n20; r21 = n21; r22 = n22; + + sign = where(any_swap, -sign, sign); + degenerate = degenerate | (best_mag < tol); + + auto pivot_safe = where(degenerate, FVal(1.0), r00); + auto f10 = r10 / pivot_safe; + auto f20 = r20 / pivot_safe; + r10 = f10; + r11 = r11 - f10 * r01; + r12 = r12 - f10 * r02; + r20 = f20; + r21 = r21 - f20 * r01; + r22 = r22 - f20 * r02; + } + + // ---- Step 1: pivot from column 1, rows {1, 2} ---- + { + auto m1 = fabs(r11), m2 = fabs(r21); + auto sel2 = (m2 > m1); + auto pivot_mag = where(sel2, m2, m1); + + auto n11 = where(sel2, r21, r11); + auto n12 = where(sel2, r22, r12); + auto n21 = where(sel2, r11, r21); + auto n22 = where(sel2, r12, r22); + r11 = n11; r12 = n12; r21 = n21; r22 = n22; + auto n10 = where(sel2, r20, r10); + auto n20 = where(sel2, r10, r20); + r10 = n10; r20 = n20; + + sign = where(sel2, -sign, sign); + degenerate = degenerate | (pivot_mag < tol); + + auto pivot_safe = where(degenerate, FVal(1.0), r11); + auto f21 = r21 / pivot_safe; + r21 = f21; + r22 = r22 - f21 * r12; + } + + degenerate = degenerate | (fabs(r22) < tol); + + auto lu_det = sign * r00 * r11 * r22; + return where(degenerate, polynomial_fallback, lu_det); +} + +// det of a symmetric 4x4 matrix via LU with partial pivoting. See lup_det3. +template +KERNELSPEC FVal lup_det4( + FVal a11, FVal a12, FVal a13, FVal a14, + FVal a22, FVal a23, FVal a24, + FVal a33, FVal a34, + FVal a44, + FVal polynomial_fallback +) { + auto r00 = a11, r01 = a12, r02 = a13, r03 = a14; + auto r10 = a12, r11 = a22, r12 = a23, r13 = a24; + auto r20 = a13, r21 = a23, r22 = a33, r23 = a34; + auto r30 = a14, r31 = a24, r32 = a34, r33 = a44; + + auto scale = max_abs(a11, a12, a13, a14, a22, a23, a24, a33, a34, a44); + auto tol = LU_TOL * scale; + auto sign = FVal(1.0); + auto degenerate = (scale == 0.); + + // ---- Step 0: pivot from column 0, rows {0,1,2,3} ---- + { + auto m0 = fabs(r00), m1 = fabs(r10), m2 = fabs(r20), m3 = fabs(r30); + auto best = IVal(0); + auto best_mag = m0; + auto pick1 = (m1 > best_mag); + best = where(pick1, IVal(1), best); + best_mag = where(pick1, m1, best_mag); + auto pick2 = (m2 > best_mag); + best = where(pick2, IVal(2), best); + best_mag = where(pick2, m2, best_mag); + auto pick3 = (m3 > best_mag); + best = where(pick3, IVal(3), best); + best_mag = where(pick3, m3, best_mag); + + auto sel1 = (best == IVal(1)); + auto sel2 = (best == IVal(2)); + auto sel3 = (best == IVal(3)); + auto any_swap = sel1 | sel2 | sel3; + + auto n00 = where(sel1, r10, where(sel2, r20, where(sel3, r30, r00))); + auto n01 = where(sel1, r11, where(sel2, r21, where(sel3, r31, r01))); + auto n02 = where(sel1, r12, where(sel2, r22, where(sel3, r32, r02))); + auto n03 = where(sel1, r13, where(sel2, r23, where(sel3, r33, r03))); + auto n10 = where(sel1, r00, r10); + auto n11 = where(sel1, r01, r11); + auto n12 = where(sel1, r02, r12); + auto n13 = where(sel1, r03, r13); + auto n20 = where(sel2, r00, r20); + auto n21 = where(sel2, r01, r21); + auto n22 = where(sel2, r02, r22); + auto n23 = where(sel2, r03, r23); + auto n30 = where(sel3, r00, r30); + auto n31 = where(sel3, r01, r31); + auto n32 = where(sel3, r02, r32); + auto n33 = where(sel3, r03, r33); + r00 = n00; r01 = n01; r02 = n02; r03 = n03; + r10 = n10; r11 = n11; r12 = n12; r13 = n13; + r20 = n20; r21 = n21; r22 = n22; r23 = n23; + r30 = n30; r31 = n31; r32 = n32; r33 = n33; + + sign = where(any_swap, -sign, sign); + degenerate = degenerate | (best_mag < tol); + + auto pivot_safe = where(degenerate, FVal(1.0), r00); + auto f10 = r10 / pivot_safe; + auto f20 = r20 / pivot_safe; + auto f30 = r30 / pivot_safe; + r10 = f10; r20 = f20; r30 = f30; + r11 = r11 - f10 * r01; r12 = r12 - f10 * r02; r13 = r13 - f10 * r03; + r21 = r21 - f20 * r01; r22 = r22 - f20 * r02; r23 = r23 - f20 * r03; + r31 = r31 - f30 * r01; r32 = r32 - f30 * r02; r33 = r33 - f30 * r03; + } + + // ---- Step 1: pivot from column 1, rows {1, 2, 3} ---- + { + auto m1 = fabs(r11), m2 = fabs(r21), m3 = fabs(r31); + auto best = IVal(1); + auto best_mag = m1; + auto pick2 = (m2 > best_mag); + best = where(pick2, IVal(2), best); + best_mag = where(pick2, m2, best_mag); + auto pick3 = (m3 > best_mag); + best = where(pick3, IVal(3), best); + best_mag = where(pick3, m3, best_mag); + + auto sel2 = (best == IVal(2)); + auto sel3 = (best == IVal(3)); + auto any_swap = sel2 | sel3; + + auto n11 = where(sel2, r21, where(sel3, r31, r11)); + auto n12 = where(sel2, r22, where(sel3, r32, r12)); + auto n13 = where(sel2, r23, where(sel3, r33, r13)); + auto n21 = where(sel2, r11, r21); + auto n22 = where(sel2, r12, r22); + auto n23 = where(sel2, r13, r23); + auto n31 = where(sel3, r11, r31); + auto n32 = where(sel3, r12, r32); + auto n33 = where(sel3, r13, r33); + auto n10 = where(sel2, r20, where(sel3, r30, r10)); + auto n20 = where(sel2, r10, r20); + auto n30 = where(sel3, r10, r30); + r10 = n10; r11 = n11; r12 = n12; r13 = n13; + r20 = n20; r21 = n21; r22 = n22; r23 = n23; + r30 = n30; r31 = n31; r32 = n32; r33 = n33; + + sign = where(any_swap, -sign, sign); + degenerate = degenerate | (best_mag < tol); + + auto pivot_safe = where(degenerate, FVal(1.0), r11); + auto f21 = r21 / pivot_safe; + auto f31 = r31 / pivot_safe; + r21 = f21; r31 = f31; + r22 = r22 - f21 * r12; r23 = r23 - f21 * r13; + r32 = r32 - f31 * r12; r33 = r33 - f31 * r13; + } + + // ---- Step 2: pivot from column 2, rows {2, 3} ---- + { + auto m2 = fabs(r22), m3 = fabs(r32); + auto sel3 = (m3 > m2); + auto pivot_mag = where(sel3, m3, m2); + + auto n22 = where(sel3, r32, r22); + auto n23 = where(sel3, r33, r23); + auto n32 = where(sel3, r22, r32); + auto n33 = where(sel3, r23, r33); + r22 = n22; r23 = n23; r32 = n32; r33 = n33; + auto n20 = where(sel3, r30, r20); + auto n21 = where(sel3, r31, r21); + auto n30 = where(sel3, r20, r30); + auto n31 = where(sel3, r21, r31); + r20 = n20; r21 = n21; r30 = n30; r31 = n31; + + sign = where(sel3, -sign, sign); + degenerate = degenerate | (pivot_mag < tol); + + auto pivot_safe = where(degenerate, FVal(1.0), r22); + auto f32 = r32 / pivot_safe; + r32 = f32; + r33 = r33 - f32 * r23; + } + + // ---- Final pivot check ---- + degenerate = degenerate | (fabs(r33) < tol); + + auto lu_det = sign * r00 * r11 * r22 * r33; + return where(degenerate, polynomial_fallback, lu_det); +} + +} // namespace kernels +} // namespace madspace diff --git a/madspace/src/kernels/threeparticle.hpp b/madspace/src/kernels/threeparticle.hpp index 3f588a5e91..e6f8af062b 100644 --- a/madspace/src/kernels/threeparticle.hpp +++ b/madspace/src/kernels/threeparticle.hpp @@ -1,6 +1,7 @@ #pragma once #include "kinematics.hpp" +#include "lup_det.hpp" namespace madspace { namespace kernels { @@ -183,9 +184,8 @@ KERNELSPEC FVal bk_V( FVal t2, FVal s12 ) { - // Determinant of the 3x3 V-matrix, - // see Eq.(11) in 10.1103/PhysRev.187.2008. - // Note: expects the absolute value of t1 + // Matrix from Byckling-Kajantie eq.(11), 10.1103/PhysRev.187.2008. + // Asymmetric (a22 = a33 = 0; a12 = a21 but a13 != a31, a23 != a32). auto a11 = 2.0 * s12; auto a12 = ma_2 + s12 - t2; auto a13 = s12 + m1_2 - m2_2; @@ -194,8 +194,18 @@ KERNELSPEC FVal bk_V( auto a31 = m0_2 + s12 - m3_2; auto a32 = m0_2 + ma_2 - mb_2; - // Computes the determinant of the 3x3 V-matrix (hard-coded because easier) - auto det = a12 * a23 * a31 + a12 * a13 * a32 - a11 * a23 * a32 - a13 * a22 * a31; + // Polynomial fallback (the previous implementation), used by the LU + // helper if the matrix's leading pivots are below the LU tolerance. + // Expansion along row 3, using a22 = a33 = 0 and a21 = a12. + auto poly_det = a31 * (a12 * a23 - a13 * a22) + + a32 * (a12 * a13 - a11 * a23); + + auto det = lup_det3_general( + a11, a12, a13, + a12, a22, a23, + a31, a32, FVal(0.0), + poly_det + ); return -det / 8.0; } @@ -212,12 +222,10 @@ KERNELSPEC FVal bk_gram4( FVal s12, FVal s23 ) { - // omputes the 4x4 Gram determinant, - // see Eq.(B6) in 10.1103/PhysRev.187.2008. - // Note: expects the absolute value of t1 + // 4x4 Gram determinant, see Eq.(B6) in 10.1103/PhysRev.187.2008. + // Note: expects the absolute value of t1. - // Get upper triangular matrix components which are non-zero - // as the Gram matrix is symmetric, i.e. (a_{ij} = a_{ji}) + // Upper-triangular entries of the symmetric Gram matrix. auto a11 = 2.0 * ma_2; auto a12 = ma_2 - t1_abs - m1_2; auto a13 = ma_2 + t2 - s12; @@ -229,14 +237,28 @@ KERNELSPEC FVal bk_gram4( auto a34 = t2 + mb_2 - m3_2; auto a44 = 2.0 * mb_2; - // Computes the determinant of the 4x4 Gram matrix (hard-coded because easier) - auto det = a14 * a23 * a14 * a23 + a13 * a24 * a13 * a24 + a12 * a34 * a12 * a34 - - a14 * a14 * a22 * a33 - a13 * a13 * a22 * a44 - a12 * a12 * a33 * a44 - - a23 * a23 * a11 * a44 - a24 * a24 * a11 * a33 - a34 * a34 * a11 * a22 + - 2 * a11 * a23 * a24 * a34 + 2 * a12 * a13 * a23 * a44 + - 2 * a12 * a14 * a24 * a33 + 2 * a13 * a14 * a22 * a34 - - 2 * a12 * a13 * a24 * a34 - 2 * a12 * a14 * a23 * a34 - - 2 * a13 * a14 * a23 * a24 + a11 * a22 * a33 * a44; + // Polynomial fallback (the previous "hard-coded because easier" expansion), + // re-organized to keep the inner 2x2 minors as their own subexpressions. + // Used by the LU helper if any pivot is below tolerance. + auto poly_det = a14 * a14 * (a23 * a23 - a22 * a33) + + a13 * a13 * (a24 * a24 - a22 * a44) + + a12 * a12 * (a34 * a34 - a33 * a44) + - a23 * a23 * a11 * a44 + - a24 * a24 * a11 * a33 + - a34 * a34 * a11 * a22 + + a11 * a22 * a33 * a44 + + 2.0 * a11 * a23 * a24 * a34 + + 2.0 * a12 * a13 * (a23 * a44 - a24 * a34) + + 2.0 * a12 * a14 * (a24 * a33 - a23 * a34) + + 2.0 * a13 * a14 * (a22 * a34 - a23 * a24); + + auto det = lup_det4( + a11, a12, a13, a14, + a22, a23, a24, + a33, a34, + a44, + poly_det + ); return det / 16.0; } @@ -252,9 +274,6 @@ KERNELSPEC FVal bk_sqrt_g3i_g3im1( FVal t2, FVal s12 ) { - // This is the squaet root of the product of the two 3x3 Gram determinants g3i and - // g3im1, as in Eq.(11) in 10.1103/PhysRev.187.2008. Note: expects the absolute - // value of t1 auto a11 = 2 * s12; auto a12 = s12 + ma_2 - t2; auto a13 = s12 + m0_2 - m3_2; @@ -265,12 +284,28 @@ KERNELSPEC FVal bk_sqrt_g3i_g3im1( auto a33 = 2 * m0_2; auto b33 = 2 * m1_2; - // Calculate the two gramm determinants g3i and g3im1 - // (ard-coded because easier) - auto g3i = a11 * a22 * a33 + 2 * a12 * a23 * a13 - a11 * a23 * a23 - - a22 * a13 * a13 - a33 * a12 * a12; - auto g3im1 = a11 * a22 * b33 + 2 * a12 * b23 * b13 - a11 * b23 * b23 - - a22 * b13 * b13 - b33 * a12 * a12; + // Polynomial fallback for each 3x3 Gram (cofactor expansion). Used by the + // LU helper independently per matrix if its pivots are below tolerance. + auto poly_g3i = a11 * (a22 * a33 - a23 * a23) + - a12 * (a12 * a33 - a13 * a23) + + a13 * (a12 * a23 - a13 * a22); + auto poly_g3im1 = a11 * (a22 * b33 - b23 * b23) + - a12 * (a12 * b33 - b13 * b23) + + b13 * (a12 * b23 - b13 * a22); + + auto g3i = lup_det3( + a11, a12, a13, + a22, a23, + a33, + poly_g3i + ); + auto g3im1 = lup_det3( + a11, a12, b13, + a22, b23, + b33, + poly_g3im1 + ); + return sqrt(g3i * g3im1) / 8.0; } @@ -567,7 +602,8 @@ KERNELSPEC void kernel_two_to_three_particle_scattering( auto p1_com = scatter_out.first; auto gram4 = bk_gram4(m0_2, ma_2, mb_2, m1_2, m2_2, m3_2, t1_abs, t2, s12, s23); auto det_2to3 = 1 / (8 * sqrt(max(-gram4, EPS2))); - auto p1_rot = rotate(p1_com, pa_com); + auto p3_p12 = boost(load_mom(p3), p_12, -1.); + auto p1_rot = rotate_two_ref(p1_com, pa_com, p3_p12); auto p1_lab = boost(p1_rot, p_12, 1.); store_mom(p1, p1_lab); for (int i = 0; i < 4; ++i) { @@ -605,7 +641,8 @@ KERNELSPEC void kernel_two_to_three_particle_scattering_inverse( auto pa_com = boost(load_mom(pa), p_12, -1.); auto p1_com = boost(load_mom(p1), p_12, -1.); - auto p1_rot = rotate_inverse(p1_com, pa_com); + auto p3_p12 = boost(load_mom(p3), p_12, -1.); + auto p1_rot = rotate_two_ref_inverse(p1_com, pa_com, p3_p12); auto m1_2 = lsquare(load_mom(p1)); auto m2_2 = lsquare(load_mom(p2)); diff --git a/madspace/src/kernels/twoparticle.hpp b/madspace/src/kernels/twoparticle.hpp index 19463ca601..39a46885d6 100644 --- a/madspace/src/kernels/twoparticle.hpp +++ b/madspace/src/kernels/twoparticle.hpp @@ -275,5 +275,76 @@ KERNELSPEC void kernel_two_to_two_particle_scattering_inverse( r_phi = phi / PI / 2. + 0.5; } +template +KERNELSPEC void kernel_double_t_scattering( + FIn r_phi, + FIn pa, + FIn pb, + FIn t1_abs, + FIn t2_abs, + FIn m1, + FOut p1, + FOut p2, + FOut det +) { + // pa + pb -> p1 + p2 with |t1| = -(pa-p1)^2, |t2| = -(pb-p1)^2, phi. + // Assumes pa, pb massless. m1 fixed; m2 derived. + // + // Canonical frame: (pa+pb) rest frame, with pa along +z. The formulas: + // E1 = (t1_abs + t2_abs + 2 m1^2) / (2 sqrt(s)) + // pz1 = (t2_abs - t1_abs) / (2 sqrt(s)) + // pt^2 = (t1_abs t2_abs + (t1_abs + t2_abs) m1^2 + m1^4 - s m1^2) / s + // Then rotate so canonical-z aligns with pa_com, boost back to lab. + FourMom p_tot; + for (int i = 0; i < 4; ++i) { + p_tot[i] = pa[i] + pb[i]; + } + auto pa_com = boost(load_mom(pa), p_tot, -1.); + auto s = lsquare(p_tot); + auto sqrts = sqrt(max(s, EPS2)); + auto m1_2 = m1 * m1; + + auto e1 = ( t1_abs + t2_abs + 2. * m1_2) / (2. * sqrts); + auto pz1 = ( t2_abs - t1_abs) / (2. * sqrts); + auto pt2_raw = + (t1_abs * t2_abs + (t1_abs + t2_abs) * m1_2 + m1_2 * m1_2 - s * m1_2) / s; + auto pt2 = max(pt2_raw, 0.); + auto pt = sqrt(pt2); + auto phi = PI * (2. * r_phi - 1.); + + FourMom p1_canon{e1, pt * cos(phi), pt * sin(phi), pz1}; + + auto p1_rot = rotate(p1_canon, pa_com); + auto p1_lab = boost(p1_rot, p_tot, 1.); + store_mom(p1, p1_lab); + for (int i = 0; i < 4; ++i) { + p2[i] = p_tot[i] - p1_lab[i]; + } + + det = PI / (2. * max(s, EPS)); +} + +template +KERNELSPEC void kernel_double_t_scattering_inverse( + FIn p1, + FIn p2, + FIn pa, + FIn pb, + FOut r_phi, + FOut det +) { + FourMom p_tot; + for (int i = 0; i < 4; ++i) { + p_tot[i] = pa[i] + pb[i]; + } + auto pa_com = boost(load_mom(pa), p_tot, -1.); + auto p1_com = boost(load_mom(p1), p_tot, -1.); + auto p1_canon = rotate_inverse(p1_com, pa_com); + auto phi = atan2(p1_canon[2], p1_canon[1]); + r_phi = phi / PI / 2. + 0.5; + auto s = lsquare(p_tot); + det = 2. * max(s, EPS) / PI; +} + } // namespace kernels } // namespace madspace diff --git a/madspace/src/phasespace/color_ordered_mapping.cpp b/madspace/src/phasespace/color_ordered_mapping.cpp new file mode 100644 index 0000000000..c405230f4f --- /dev/null +++ b/madspace/src/phasespace/color_ordered_mapping.cpp @@ -0,0 +1,554 @@ +#include "madspace/phasespace/color_ordered_mapping.hpp" + +#include +#include + +#include "madspace/util.hpp" + +using namespace madspace; + +namespace { + +// Cyclically rotate the color order so that particle 0 is first, then +// split the rest into two sets: particles strictly between 0 and 1 in the +// cyclic order go into set1, particles after 1 (wrapping around) go into +// set2. +std::pair, std::vector> +split_sets_from_color_order(const std::vector& color_order) { + std::size_t n = color_order.size(); + if (n < 4) { + throw std::invalid_argument( + "ColorOrderedMapping requires at least 4 particles (2 beams + 2 outgoing)" + ); + } + auto it = std::find(color_order.begin(), color_order.end(), 0u); + if (it == color_order.end()) { + throw std::invalid_argument("color_order must contain particle 0"); + } + std::size_t i0 = std::distance(color_order.begin(), it); + std::vector rotated; + rotated.reserve(n); + for (std::size_t k = 0; k < n; ++k) { + rotated.push_back(color_order[(i0 + k) % n]); + } + auto it1 = std::find(rotated.begin(), rotated.end(), 1u); + if (it1 == rotated.end()) { + throw std::invalid_argument("color_order must contain particle 1"); + } + std::size_t i1 = std::distance(rotated.begin(), it1); + std::vector set1, set2; + for (std::size_t k = 1; k < i1; ++k) { + std::size_t p = rotated[k]; + if (p <= 1) throw std::invalid_argument("invalid color_order"); + set1.push_back(p - 2); + } + for (std::size_t k = i1 + 1; k < n; ++k) { + std::size_t p = rotated[k]; + if (p <= 1) throw std::invalid_argument("invalid color_order"); + set2.push_back(p - 2); + } + if (set1.empty() || set2.empty()) { + throw std::invalid_argument( + "ColorOrderedMapping requires both sets to be non-empty " + "(particles 0 and 1 must not be adjacent in the color order)" + ); + } + return {set1, set2}; +} + +std::size_t n_intermediate_masses_for_set_size(std::size_t k) { + return (k >= 2) ? (k - 2) : 0; +} + +std::size_t n_block_randoms_for_set_size(std::size_t k) { + if (k <= 1) return 0; + // First peel: 2->2 LAB (2 randoms); each subsequent peel: 2->3 (3 randoms). + return 2 + 3 * (k - 2); +} + +} // namespace + +ColorOrderedMapping::ColorOrderedMapping( + const std::vector& color_order, + double t_invariant_power, + double s_invariant_power +) : + Mapping( + "ColorOrderedMapping", + [&] { + auto [s1, s2] = split_sets_from_color_order(color_order); + bool use_double_t = (s1.size() == 1) != (s2.size() == 1); + // DoubleT branch: 0 set-masses + 3 central randoms. + // Standard branch: 1 set-mass per multi-particle side + 2 central randoms. + std::size_t n_set_masses = use_double_t + ? 0u + : (s1.size() >= 2 ? 1u : 0u) + (s2.size() >= 2 ? 1u : 0u); + std::size_t n_central = use_double_t ? 3u : 2u; + std::size_t n_intermediate_masses = + n_intermediate_masses_for_set_size(s1.size()) + + n_intermediate_masses_for_set_size(s2.size()); + std::size_t n_walk = + n_block_randoms_for_set_size(s1.size()) + + n_block_randoms_for_set_size(s2.size()); + std::size_t total = n_set_masses + n_intermediate_masses + + n_central + n_walk; + NamedVector input_types; + for (std::size_t i = 0; i < total; ++i) { + input_types.push_back(std::format("random{}", i), batch_float); + } + return input_types; + }(), + [&] { + std::size_t n_out = color_order.size() - 2; + NamedVector output_types; + for (std::size_t i = 0; i < n_out + 2; ++i) { + output_types.push_back(std::format("momentum{}", i), batch_four_vec); + } + return output_types; + }(), + [&] { + std::size_t n_out = color_order.size() - 2; + NamedVector cond_types{{"com_energy", batch_float}}; + for (std::size_t i = 0; i < n_out; ++i) { + cond_types.push_back(std::format("mass{}", i), batch_float); + } + return cond_types; + }() + ), + _n_out(color_order.size() - 2), + _com_scattering(true, t_invariant_power), + _lab_scattering(false, t_invariant_power), + _two_to_three(t_invariant_power, 0., 0., s_invariant_power, 0., 0.), + _double_t(t_invariant_power, 0., 0., t_invariant_power, 0., 0.) { + auto [s1, s2] = split_sets_from_color_order(color_order); + _set1 = s1; + _set2 = s2; + _use_double_t = (s1.size() == 1) != (s2.size() == 1); + std::size_t n_set_masses = _use_double_t + ? 0u + : (s1.size() >= 2 ? 1u : 0u) + (s2.size() >= 2 ? 1u : 0u); + std::size_t n_central = _use_double_t ? 3u : 2u; + std::size_t n_intermediate_masses = + n_intermediate_masses_for_set_size(s1.size()) + + n_intermediate_masses_for_set_size(s2.size()); + std::size_t n_walk = + n_block_randoms_for_set_size(s1.size()) + + n_block_randoms_for_set_size(s2.size()); + _random_dim = n_set_masses + n_intermediate_masses + n_central + n_walk; +} + +Mapping::Result ColorOrderedMapping::build_forward_impl( + FunctionBuilder& fb, + const NamedVector& inputs, + const NamedVector& conditions +) const { + Value e_cm = conditions.at(0); + ValueVec m_out(conditions.begin() + 1, conditions.end()); + auto r = inputs.begin(); + auto next_random = [&]() { return *(r++); }; + ValueVec dets; + + // Phase 1a: pre-sample set composite masses. + // DoubleT branch defers the multi-particle side's mass until after Phase 2. + Value m_set1, m_set2; + auto masses_of = [&](const std::vector& idxs) { + ValueVec v; v.reserve(idxs.size()); + for (auto i : idxs) v.push_back(m_out.at(i)); + return fb.sum(v); + }; + Value mass_sum_set1 = masses_of(_set1); + Value mass_sum_set2 = masses_of(_set2); + if (_use_double_t) { + if (_set1.size() == 1) m_set1 = m_out.at(_set1[0]); + if (_set2.size() == 1) m_set2 = m_out.at(_set2[0]); + } else { + if (_set1.size() >= 2) { + auto s_min = fb.square(mass_sum_set1); + auto s_max = fb.square(fb.sub(e_cm, mass_sum_set2)); + auto res = _uniform_invariant.build_forward(fb, {next_random()}, {s_min, s_max}); + m_set1 = fb.sqrt(res["invariant"]); + dets.push_back(res["det"]); + } else { + m_set1 = m_out.at(_set1[0]); + } + if (_set2.size() >= 2) { + auto s_min = fb.square(mass_sum_set2); + auto s_max = fb.square(fb.sub(e_cm, m_set1)); + auto res = _uniform_invariant.build_forward(fb, {next_random()}, {s_min, s_max}); + m_set2 = fb.sqrt(res["invariant"]); + dets.push_back(res["det"]); + } else { + m_set2 = m_out.at(_set2[0]); + } + } + + // Phase 2: central block producing (P_set1, P_set2). + auto [pa, pb] = fb.com_p_in(e_cm); + Value P_set1, P_set2; + if (_use_double_t) { + // DoubleT: pa, pb -> (single, recoil). The single particle's mass + // is m_single (fixed); the recoil mass must be >= mir_min (sum of + // recoil-side outgoing masses). + bool single_is_set1 = (_set1.size() == 1); + Value m_single = single_is_set1 ? m_set1 : m_set2; + Value mir_min = single_is_set1 ? mass_sum_set2 : mass_sum_set1; + auto central = _double_t.build_forward( + fb, + {next_random(), next_random(), next_random()}, + {pa, pb, m_single, mir_min} + ); + Value p_single = central.at(0); + Value p_recoil = central.at(1); + dets.push_back(central["det"]); + if (single_is_set1) { + P_set1 = p_single; + P_set2 = p_recoil; + // Derive multi-side m_set from the recoil momentum for Phase 1b. + nested_vector2 factors_recoil{{1.0}}; + auto m2 = fb.unstack( + fb.invariants_from_momenta(fb.stack({p_recoil}), factors_recoil) + ).at(0); + m_set2 = fb.sqrt(m2); + } else { + P_set2 = p_single; + P_set1 = p_recoil; + nested_vector2 factors_recoil{{1.0}}; + auto m2 = fb.unstack( + fb.invariants_from_momenta(fb.stack({p_recoil}), factors_recoil) + ).at(0); + m_set1 = fb.sqrt(m2); + } + } else { + auto central = _com_scattering.build_forward( + fb, + {next_random(), next_random(), m_set1, m_set2}, + {pa, pb} + ); + P_set1 = central.at(0); + P_set2 = central.at(1); + dets.push_back(central["det"]); + } + + // R_b for each walk: the beam minus what the central block emitted on the other side. + Value R_b_for_set1 = fb.sub(pb, P_set2); + Value R_b_for_set2 = fb.sub(pa, P_set1); + + // Phase 1b: pre-sample intermediate rest masses for each walk. + // For each step j in 0..k-3 (last is the residual mass(s[k-1])): + // m_min = sum masses [set[j+1]..set[k-1]] + // m_max = m_rest_j - mass[set[j]], m_rest_0 = m_set + // Sampling from [m_min, m_set - sum_peeled] independently would + // over-count by (k-2)! by ignoring the monotonic ordering. + // Sampled after the central block: DoubleT only fixes m_set then. + auto sample_intermediate_masses = + [&](const std::vector& s, Value m_set) -> ValueVec { + std::size_t k = s.size(); + ValueVec res; + if (k <= 2) { + return res; + } + Value prev_mass = m_set; + for (std::size_t j = 0; j < k - 2; ++j) { + Value m_min = m_out.at(s[j + 1]); + for (std::size_t i = j + 2; i < k; ++i) { + m_min = fb.add(m_min, m_out.at(s[i])); + } + auto s_min = fb.square(m_min); + auto s_max = fb.square(fb.sub(prev_mass, m_out.at(s[j]))); + auto r = _uniform_invariant.build_forward(fb, {next_random()}, {s_min, s_max}); + Value m_rest = fb.sqrt(r["invariant"]); + res.push_back(m_rest); + dets.push_back(r["det"]); + prev_mass = m_rest; + } + return res; + }; + + ValueVec rest_masses_set1 = sample_intermediate_masses(_set1, m_set1); + ValueVec rest_masses_set2 = sample_intermediate_masses(_set2, m_set2); + + // Phase 3: peel-off walks + ValueVec p_out(_n_out); + + auto walk = [&](const std::vector& s, + Value P_set, + Value R_b, + const ValueVec& rest_masses) { + std::size_t k = s.size(); + if (k == 1) { + p_out[s[0]] = P_set; + return; + } + Value R_a = fb.sub(P_set, R_b); + Value im1; + bool first = true; + for (std::size_t j = 0; j < k - 1; ++j) { + // New-rest mass: intermediate (rest_masses[j]) or final residual (j == k-2). + Value m_rest = (j < k - 2) ? rest_masses[j] : m_out.at(s[k - 1]); + Value m_peel = m_out.at(s[j]); + // Block convention: pa-side carries the chain (mass m_rest), pb-side + // carries the peeled particle (mass m_peel). R_a is the active leg + // and gets decremented by peeled; R_b is constant. + // + // 2->3 kernel internally subtracts p_3 = im1 from pa+pb. R_a already + // has im1 subtracted from our previous step, so we pass pb = R_a + im1 + // to recover p_12 = R_b + R_a inside the kernel. + if (first) { + auto ks = _lab_scattering.build_forward( + fb, + {next_random(), next_random(), m_rest, m_peel}, + {R_b, R_a} + ); + Value peeled = ks.at(1); + p_out[s[j]] = peeled; + R_a = fb.sub(R_a, peeled); + im1 = peeled; + dets.push_back(ks["det"]); + first = false; + } else { + Value pb_for_block = fb.add(R_a, im1); + auto ks = _two_to_three.build_forward( + fb, + {next_random(), next_random(), next_random(), m_rest, m_peel}, + {R_b, pb_for_block, im1} + ); + Value peeled = ks.at(1); + p_out[s[j]] = peeled; + R_a = fb.sub(R_a, peeled); + im1 = peeled; + dets.push_back(ks["det"]); + } + } + // Last particle = R_a + R_b (= what remains after all explicit peels). + p_out[s[k - 1]] = fb.add(R_a, R_b); + }; + + walk(_set1, P_set1, R_b_for_set1, rest_masses_set1); + walk(_set2, P_set2, R_b_for_set2, rest_masses_set2); + + // Assemble outputs: momentum0, momentum1 = beams; momentum_{2+i} = outgoing i. + ValueVec p_ext; + p_ext.reserve(_n_out + 2); + p_ext.push_back(pa); + p_ext.push_back(pb); + for (std::size_t i = 0; i < _n_out; ++i) { + p_ext.push_back(p_out[i]); + } + + return {{output_types().keys(), p_ext}, fb.product(dets)}; +} + +Mapping::Result ColorOrderedMapping::build_inverse_impl( + FunctionBuilder& fb, + const NamedVector& inputs, + const NamedVector& conditions +) const { + Value e_cm = conditions.at(0); + ValueVec m_out(conditions.begin() + 1, conditions.end()); + ValueVec random_out; + ValueVec dets; + + Value pa = inputs.at(0); + Value pb = inputs.at(1); + auto p_outgoing = [&](std::size_t i) { return inputs.at(2 + i); }; + + // Factor matrix: extract all needed invariant masses in one fb call. + // Layout: m^2(P_set1), m^2(P_set2), then per-set rest-system masses. + // idx_* below record where each band starts. + nested_vector2 invariant_factors; + auto n_inputs = _n_out + 2; + auto factor_set = [&](const std::vector& idxs) { + std::vector row(n_inputs, 0.0); + for (auto idx : idxs) row.at(idx + 2) = 1.0; + return row; + }; + int idx_m2_set1 = -1, idx_m2_set2 = -1; + int idx_rest_set1_start = -1, idx_rest_set2_start = -1; + if (_set1.size() >= 2) { + idx_m2_set1 = invariant_factors.size(); + invariant_factors.push_back(factor_set(_set1)); + } + if (_set2.size() >= 2) { + idx_m2_set2 = invariant_factors.size(); + invariant_factors.push_back(factor_set(_set2)); + } + if (_set1.size() >= 3) { + idx_rest_set1_start = invariant_factors.size(); + for (std::size_t j = 0; j < _set1.size() - 2; ++j) { + std::vector sub(_set1.begin() + j + 1, _set1.end()); + invariant_factors.push_back(factor_set(sub)); + } + } + if (_set2.size() >= 3) { + idx_rest_set2_start = invariant_factors.size(); + for (std::size_t j = 0; j < _set2.size() - 2; ++j) { + std::vector sub(_set2.begin() + j + 1, _set2.end()); + invariant_factors.push_back(factor_set(sub)); + } + } + ValueVec invariants; + if (!invariant_factors.empty()) { + invariants = fb.unstack( + fb.invariants_from_momenta(fb.stack(inputs.values()), invariant_factors) + ); + } + + auto sum_momenta = [&](const std::vector& s) -> Value { + Value p = p_outgoing(s[0]); + for (std::size_t k = 1; k < s.size(); ++k) { + p = fb.add(p, p_outgoing(s[k])); + } + return p; + }; + Value P_set1 = sum_momenta(_set1); + Value P_set2 = sum_momenta(_set2); + + // Phase 1a inverse: recover set-mass randoms (none for DoubleT). + // m_set values come from invariants_from_momenta for multi-particle sides. + Value m_set1, m_set2; + auto masses_of = [&](const std::vector& idxs) { + ValueVec v; v.reserve(idxs.size()); + for (auto i : idxs) v.push_back(m_out.at(i)); + return fb.sum(v); + }; + Value mass_sum_set1 = masses_of(_set1); + Value mass_sum_set2 = masses_of(_set2); + if (_use_double_t) { + if (_set1.size() == 1) { + m_set1 = m_out.at(_set1[0]); + m_set2 = fb.sqrt(invariants.at(idx_m2_set2)); + } else { + m_set2 = m_out.at(_set2[0]); + m_set1 = fb.sqrt(invariants.at(idx_m2_set1)); + } + } else { + if (_set1.size() >= 2) { + auto s_min = fb.square(mass_sum_set1); + auto s_max = fb.square(fb.sub(e_cm, mass_sum_set2)); + auto m2_set1 = invariants.at(idx_m2_set1); + auto res = _uniform_invariant.build_inverse(fb, {m2_set1}, {s_min, s_max}); + random_out.push_back(res["random"]); + dets.push_back(res["det"]); + m_set1 = fb.sqrt(m2_set1); + } else { + m_set1 = m_out.at(_set1[0]); + } + if (_set2.size() >= 2) { + auto s_min = fb.square(mass_sum_set2); + auto s_max = fb.square(fb.sub(e_cm, m_set1)); + auto m2_set2 = invariants.at(idx_m2_set2); + auto res = _uniform_invariant.build_inverse(fb, {m2_set2}, {s_min, s_max}); + random_out.push_back(res["random"]); + dets.push_back(res["det"]); + m_set2 = fb.sqrt(m2_set2); + } else { + m_set2 = m_out.at(_set2[0]); + } + } + + // Phase 2 inverse: central block. Comes before Phase 1b because the + // forward emits central randoms first (so multi-side m_set is known by 1b). + if (_use_double_t) { + bool single_is_set1 = (_set1.size() == 1); + Value p_single = single_is_set1 ? P_set1 : P_set2; + Value p_recoil = single_is_set1 ? P_set2 : P_set1; + Value m_single = single_is_set1 ? m_set1 : m_set2; + Value mir_min = single_is_set1 ? mass_sum_set2 : mass_sum_set1; + auto central = _double_t.build_inverse( + fb, + {p_single, p_recoil}, + {pa, pb, m_single, mir_min} + ); + random_out.push_back(central.at(0)); + random_out.push_back(central.at(1)); + random_out.push_back(central.at(2)); + dets.push_back(central["det"]); + } else { + auto central = _com_scattering.build_inverse( + fb, + {P_set1, P_set2}, + {pa, pb} + ); + random_out.push_back(central.at(0)); + random_out.push_back(central.at(1)); + dets.push_back(central["det"]); + } + + // Phase 1b inverse: recover intermediate rest-mass randoms + auto recover_intermediate_masses = + [&](const std::vector& s, Value m_set, int idx_start) { + std::size_t k = s.size(); + if (k <= 2) return; + // Mirror the forward chaining: prev_mass starts at m_set, then becomes + // the (sqrt of the) actual recovered m2 at each step. + Value prev_mass = m_set; + for (std::size_t j = 0; j < k - 2; ++j) { + Value m_min = m_out.at(s[j + 1]); + for (std::size_t i = j + 2; i < k; ++i) { + m_min = fb.add(m_min, m_out.at(s[i])); + } + auto s_min = fb.square(m_min); + auto s_max = fb.square(fb.sub(prev_mass, m_out.at(s[j]))); + auto m2 = invariants.at(idx_start + j); + auto res = _uniform_invariant.build_inverse(fb, {m2}, {s_min, s_max}); + random_out.push_back(res["random"]); + dets.push_back(res["det"]); + prev_mass = fb.sqrt(m2); + } + }; + + recover_intermediate_masses(_set1, m_set1, idx_rest_set1_start); + recover_intermediate_masses(_set2, m_set2, idx_rest_set2_start); + + Value R_b_for_set1 = fb.sub(pb, P_set2); + Value R_b_for_set2 = fb.sub(pa, P_set1); + + // Phase 3 inverse: peel-off walks + auto walk_inverse = [&](const std::vector& s, + Value P_set, + Value R_b) { + std::size_t k = s.size(); + if (k == 1) return; + Value R_a = fb.sub(P_set, R_b); + Value im1; + bool first = true; + for (std::size_t j = 0; j < k - 1; ++j) { + Value peeled = p_outgoing(s[j]); + // p1_out is the chain carrier (mass m_rest); by conservation + // p1_out = R_a + R_b - peeled. + Value p1_out = fb.sub(fb.add(R_a, R_b), peeled); + if (first) { + auto rs = _lab_scattering.build_inverse( + fb, + {p1_out, peeled}, + {R_b, R_a} + ); + random_out.push_back(rs.at(0)); + random_out.push_back(rs.at(1)); + dets.push_back(rs["det"]); + first = false; + } else { + // Mirror the forward's pb restoration: the 2->3 kernel + // internally subtracts p_3 = im1, so we pass pb = R_a + im1 + // to get p_12 = R_b + R_a (the remaining-to-produce system). + Value pb_for_block = fb.add(R_a, im1); + auto rs = _two_to_three.build_inverse( + fb, + {p1_out, peeled}, + {R_b, pb_for_block, im1} + ); + random_out.push_back(rs.at(0)); + random_out.push_back(rs.at(1)); + random_out.push_back(rs.at(2)); + dets.push_back(rs["det"]); + } + R_a = fb.sub(R_a, peeled); + im1 = peeled; + } + }; + + walk_inverse(_set1, P_set1, R_b_for_set1); + walk_inverse(_set2, P_set2, R_b_for_set2); + + return {{input_types().keys(), random_out}, fb.product(dets)}; +} diff --git a/madspace/src/phasespace/two_particle.cpp b/madspace/src/phasespace/two_particle.cpp index 634997c750..fd85ae8601 100644 --- a/madspace/src/phasespace/two_particle.cpp +++ b/madspace/src/phasespace/two_particle.cpp @@ -125,3 +125,81 @@ Mapping::Result TwoToTwoParticleScattering::build_inverse_impl( fb.mul(t_result["det"], det_scatter) }; } + +DoubleT::DoubleT( + double t1_invariant_power, double t1_mass, double t1_width, + double t2_invariant_power, double t2_mass, double t2_width +) : + Mapping( + "DoubleT", + {{"random_phi", batch_float}, + {"random_t1", batch_float}, + {"random_t2", batch_float}}, + {{"momentum1", batch_four_vec}, {"momentum2", batch_four_vec}}, + {{"momentum_in1", batch_four_vec}, + {"momentum_in2", batch_four_vec}, + {"mass1", batch_float}, + {"mass_rest_min", batch_float}} + ), + _t1_invariant(t1_invariant_power, t1_mass, t1_width), + _t2_invariant(t2_invariant_power, t2_mass, t2_width) {} + +Mapping::Result DoubleT::build_forward_impl( + FunctionBuilder& fb, + const NamedVector& inputs, + const NamedVector& conditions +) const { + auto r_phi = inputs.at(0), r_t1 = inputs.at(1), r_t2 = inputs.at(2); + auto p_in1 = conditions.at(0), p_in2 = conditions.at(1); + auto m1 = conditions.at(2); + auto mir_min = conditions.at(3); + + auto [t1_min, t1_max] = fb.t1_inv_min_max_doublet(p_in1, p_in2, m1, mir_min); + auto t1_result = _t1_invariant.build_forward(fb, {r_t1}, {t1_min, t1_max}); + + auto [t2_min, t2_max] = + fb.t2_inv_min_max_doublet(p_in1, p_in2, m1, mir_min, t1_result["invariant"]); + auto t2_result = _t2_invariant.build_forward(fb, {r_t2}, {t2_min, t2_max}); + + auto [p1, p2, det_scatter] = fb.double_t_scattering( + r_phi, p_in1, p_in2, + t1_result["invariant"], t2_result["invariant"], m1 + ); + + auto det_inv = fb.mul(t1_result["det"], t2_result["det"]); + return { + {{"momentum1", p1}, {"momentum2", p2}}, + fb.mul(det_inv, det_scatter) + }; +} + +Mapping::Result DoubleT::build_inverse_impl( + FunctionBuilder& fb, + const NamedVector& inputs, + const NamedVector& conditions +) const { + auto p1 = inputs.at(0), p2 = inputs.at(1); + auto p_in1 = conditions.at(0), p_in2 = conditions.at(1); + auto m1 = conditions.at(2); + auto mir_min = conditions.at(3); + + auto [r_phi, det_scatter] = + fb.double_t_scattering_inverse(p1, p2, p_in1, p_in2); + + auto [t1_abs, t1_min, t1_max] = + fb.t1_inv_value_and_min_max_doublet(p_in1, p_in2, p1, m1, mir_min); + auto t1_result = _t1_invariant.build_inverse(fb, {t1_abs}, {t1_min, t1_max}); + + auto [t2_abs, t2_min, t2_max] = fb.t2_inv_value_and_min_max_doublet( + p_in1, p_in2, p1, m1, mir_min, t1_abs + ); + auto t2_result = _t2_invariant.build_inverse(fb, {t2_abs}, {t2_min, t2_max}); + + auto det_inv = fb.mul(t1_result["det"], t2_result["det"]); + return { + {{"random_phi", r_phi}, + {"random_t1", t1_result["random"]}, + {"random_t2", t2_result["random"]}}, + fb.mul(det_inv, det_scatter) + }; +} diff --git a/madspace/src/python/instruction_set.hpp b/madspace/src/python/instruction_set.hpp index ac48f4d7aa..5f3509a19e 100644 --- a/madspace/src/python/instruction_set.hpp +++ b/madspace/src/python/instruction_set.hpp @@ -34,8 +34,8 @@ void add_instructions(py::classh& fb) { fb.def("div", &FunctionBuilder::div, py::arg("in1"), py::arg("in2")); fb.def("reduce_sum", &FunctionBuilder::reduce_sum, py::arg("in1")); fb.def("reduce_sum_vector", &FunctionBuilder::reduce_sum_vector, py::arg("in1")); - fb.def("batch_reduce_mean", &FunctionBuilder::batch_reduce_mean, py::arg("in")); - fb.def("batch_reduce_mean_keepdim", &FunctionBuilder::batch_reduce_mean_keepdim, py::arg("in")); + fb.def("batch_reduce_mean", &FunctionBuilder::batch_reduce_mean, py::arg("in1")); + fb.def("batch_reduce_mean_keepdim", &FunctionBuilder::batch_reduce_mean_keepdim, py::arg("in1")); fb.def("reduce_product", &FunctionBuilder::reduce_product, py::arg("in1")); fb.def("sqrt", &FunctionBuilder::sqrt, py::arg("in1")); fb.def("square", &FunctionBuilder::square, py::arg("in1")); @@ -74,12 +74,18 @@ void add_instructions(py::classh& fb) { fb.def("two_to_two_particle_scattering_inverse", &FunctionBuilder::two_to_two_particle_scattering_inverse, py::arg("p1"), py::arg("p2"), py::arg("pa"), py::arg("pb")); fb.def("two_to_three_particle_scattering", &FunctionBuilder::two_to_three_particle_scattering, py::arg("phi_choice"), py::arg("pa"), py::arg("pb"), py::arg("p3"), py::arg("s23"), py::arg("t1_abs"), py::arg("m1"), py::arg("m2")); fb.def("two_to_three_particle_scattering_inverse", &FunctionBuilder::two_to_three_particle_scattering_inverse, py::arg("p1"), py::arg("p2"), py::arg("p3"), py::arg("pa"), py::arg("pb"), py::arg("t1_abs"), py::arg("s23")); + fb.def("double_t_scattering", &FunctionBuilder::double_t_scattering, py::arg("r_phi"), py::arg("pa"), py::arg("pb"), py::arg("t1_abs"), py::arg("t2_abs"), py::arg("m1")); + fb.def("double_t_scattering_inverse", &FunctionBuilder::double_t_scattering_inverse, py::arg("p1"), py::arg("p2"), py::arg("pa"), py::arg("pb")); fb.def("three_body_decay_com", &FunctionBuilder::three_body_decay_com, py::arg("r_e1"), py::arg("r_e2"), py::arg("r_phi"), py::arg("r_cos_theta"), py::arg("r_beta"), py::arg("m0"), py::arg("m1"), py::arg("m2"), py::arg("m3")); fb.def("three_body_decay_com_inverse", &FunctionBuilder::three_body_decay_com_inverse, py::arg("p1"), py::arg("p2"), py::arg("p3")); fb.def("three_body_decay", &FunctionBuilder::three_body_decay, py::arg("r_e1"), py::arg("r_e2"), py::arg("r_phi"), py::arg("r_cos_theta"), py::arg("r_beta"), py::arg("m0"), py::arg("m1"), py::arg("m2"), py::arg("m3"), py::arg("p0")); fb.def("three_body_decay_inverse", &FunctionBuilder::three_body_decay_inverse, py::arg("p1"), py::arg("p2"), py::arg("p3")); fb.def("t_inv_min_max", &FunctionBuilder::t_inv_min_max, py::arg("pa"), py::arg("pb"), py::arg("m1"), py::arg("m2")); fb.def("t_inv_value_and_min_max", &FunctionBuilder::t_inv_value_and_min_max, py::arg("pa"), py::arg("pb"), py::arg("p1"), py::arg("p2")); + fb.def("t1_inv_min_max_doublet", &FunctionBuilder::t1_inv_min_max_doublet, py::arg("pa"), py::arg("pb"), py::arg("m1"), py::arg("mir_min")); + fb.def("t1_inv_value_and_min_max_doublet", &FunctionBuilder::t1_inv_value_and_min_max_doublet, py::arg("pa"), py::arg("pb"), py::arg("p1"), py::arg("m1"), py::arg("mir_min")); + fb.def("t2_inv_min_max_doublet", &FunctionBuilder::t2_inv_min_max_doublet, py::arg("pa"), py::arg("pb"), py::arg("m1"), py::arg("mir_min"), py::arg("t1_abs")); + fb.def("t2_inv_value_and_min_max_doublet", &FunctionBuilder::t2_inv_value_and_min_max_doublet, py::arg("pa"), py::arg("pb"), py::arg("p1"), py::arg("m1"), py::arg("mir_min"), py::arg("t1_abs")); fb.def("s23_min_max", &FunctionBuilder::s23_min_max, py::arg("pa"), py::arg("pb"), py::arg("p3"), py::arg("t1_abs"), py::arg("m1"), py::arg("m2")); fb.def("s23_value_and_min_max", &FunctionBuilder::s23_value_and_min_max, py::arg("pa"), py::arg("pb"), py::arg("p3"), py::arg("t1_abs"), py::arg("p1"), py::arg("p2")); fb.def("invariants_from_momenta", &FunctionBuilder::invariants_from_momenta, py::arg("p_ext"), py::arg("factors")); diff --git a/madspace/src/python/madspace.cpp b/madspace/src/python/madspace.cpp index 8f5d080395..c9cf692596 100644 --- a/madspace/src/python/madspace.cpp +++ b/madspace/src/python/madspace.cpp @@ -438,6 +438,17 @@ PYBIND11_MODULE(_madspace_py, m) { py::arg("width") = 0. ); + py::classh(m, "DoubleT") + .def( + py::init(), + py::arg("t1_invariant_power") = 0., + py::arg("t1_mass") = 0., + py::arg("t1_width") = 0., + py::arg("t2_invariant_power") = 0., + py::arg("t2_mass") = 0., + py::arg("t2_width") = 0. + ); + py::classh(m, "ThreeBodyDecay") .def(py::init(), py::arg("com")); @@ -474,7 +485,17 @@ PYBIND11_MODULE(_madspace_py, m) { py::init, double>(), py::arg("integration_order"), py::arg("invariant_power") = 0. - ); + ) + .def("random_dim", &TPropagatorMapping::random_dim); + + py::classh(m, "ColorOrderedMapping") + .def( + py::init, double, double>(), + py::arg("color_order"), + py::arg("t_invariant_power") = 0., + py::arg("s_invariant_power") = 0. + ) + .def("random_dim", &ColorOrderedMapping::random_dim); py::classh(m, "VegasHistogram") .def( @@ -1411,81 +1432,6 @@ PYBIND11_MODULE(_madspace_py, m) { .def_readonly("bin_values", &Histogram::bin_values) .def_readonly("bin_errors", &Histogram::bin_errors); - py::classh(m, "ChannelEventGenerator") - .def_static( - "load", - &ChannelEventGenerator::load, - py::arg("channel_file"), - py::arg("contexts"), - py::arg("event_file"), - py::arg("weight_file"), - py::arg("config") - ) - .def( - py::init< - const std::vector&, - const Integrand&, - const std::string&, - const std::string&, - const GeneratorConfig&, - std::size_t, - const std::string&, - const std::optional&>(), - py::arg("contexts"), - py::arg("integrand"), - py::arg("event_file"), - py::arg("weight_file"), - py::arg("config"), - py::arg("subprocess_index"), - py::arg("name"), - py::arg("histograms") - ) - .def_readonly_static("integrand_flags", &ChannelEventGenerator::integrand_flags) - .def("status", &ChannelEventGenerator::status) - .def("save", &ChannelEventGenerator::save, py::arg("save")); - - py::classh(m, "EventGenerator") - .def_readonly_static("default_config", &EventGenerator::default_config) - .def( - py::init< - const std::vector&, - const std::vector>&, - const std::string&, - const GeneratorConfig&>(), - py::arg("contexts"), - py::arg("channels"), - py::arg("status_file") = "", - py::arg_v( - "config", - EventGenerator::default_config, - "EventGenerator.default_config" - ) - ) - .def("survey", &EventGenerator::survey) - .def("generate", &EventGenerator::generate) - .def( - "combine_to_compact_npy", - &EventGenerator::combine_to_compact_npy, - py::arg("file_name") - ) - .def( - "combine_to_lhe_npy", - &EventGenerator::combine_to_lhe_npy, - py::arg("file_name"), - py::arg("lhe_completer") - ) - .def( - "combine_to_lhe", - &EventGenerator::combine_to_lhe, - py::arg("file_name"), - py::arg("lhe_completer") - ) - .def("status", &EventGenerator::status) - .def("channel_status", &EventGenerator::channel_status) - .def("histograms", &EventGenerator::histograms) - .def("used_globals", &EventGenerator::used_globals) - .def("channels", &EventGenerator::channels); - py::classh(m, "LHEHeader") .def( py::init(), @@ -1705,6 +1651,81 @@ PYBIND11_MODULE(_madspace_py, m) { .def("print_update", &PrettyBox::print_update) .def_property_readonly("line_count", &PrettyBox::line_count); + py::classh(m, "ChannelEventGenerator") + .def_static( + "load", + &ChannelEventGenerator::load, + py::arg("channel_file"), + py::arg("contexts"), + py::arg("event_file"), + py::arg("weight_file"), + py::arg("config") + ) + .def( + py::init< + const std::vector&, + const Integrand&, + const std::string&, + const std::string&, + const GeneratorConfig&, + std::size_t, + const std::string&, + const std::optional&>(), + py::arg("contexts"), + py::arg("integrand"), + py::arg("event_file"), + py::arg("weight_file"), + py::arg("config"), + py::arg("subprocess_index"), + py::arg("name"), + py::arg("histograms") + ) + .def_readonly_static("integrand_flags", &ChannelEventGenerator::integrand_flags) + .def("status", &ChannelEventGenerator::status) + .def("save", &ChannelEventGenerator::save, py::arg("save")); + + py::classh(m, "EventGenerator") + .def_readonly_static("default_config", &EventGenerator::default_config) + .def( + py::init< + const std::vector&, + const std::vector>&, + const std::string&, + const GeneratorConfig&>(), + py::arg("contexts"), + py::arg("channels"), + py::arg("status_file") = "", + py::arg_v( + "config", + EventGenerator::default_config, + "EventGenerator.default_config" + ) + ) + .def("survey", &EventGenerator::survey) + .def("generate", &EventGenerator::generate) + .def( + "combine_to_compact_npy", + &EventGenerator::combine_to_compact_npy, + py::arg("file_name") + ) + .def( + "combine_to_lhe_npy", + &EventGenerator::combine_to_lhe_npy, + py::arg("file_name"), + py::arg("lhe_completer") + ) + .def( + "combine_to_lhe", + &EventGenerator::combine_to_lhe, + py::arg("file_name"), + py::arg("lhe_completer") + ) + .def("status", &EventGenerator::status) + .def("channel_status", &EventGenerator::channel_status) + .def("histograms", &EventGenerator::histograms) + .def("used_globals", &EventGenerator::used_globals) + .def("channels", &EventGenerator::channels); + py::classh logger(m, "Logger"); add_enum( logger, diff --git a/madspace/tests/test_2to3_scattering.py b/madspace/tests/test_2to3_scattering.py index 6136bb3439..dfa52bed62 100644 --- a/madspace/tests/test_2to3_scattering.py +++ b/madspace/tests/test_2to3_scattering.py @@ -40,7 +40,6 @@ def boost(k: np.ndarray, p_boost: np.ndarray, inverse: bool = False) -> np.ndarr # Fixtures # ---------------------------- -# N = 50_000 # keep this moderate for CI speed; bump locally for tighter stats N = 10_000 # keep this moderate for CI speed; bump locally for tighter stats InputPoint = namedtuple( @@ -186,30 +185,28 @@ def test_momentum_conservation(input_points): assert p1 + p2 == approx(input_points.p12) -# TODO: Find bug in inverse mapping -# def test_inverse(input_points): -# mapping = ms.TwoToThreeParticleScattering() - -# inputs = [ -# input_points.r_choice, -# input_points.r_s23, -# input_points.r_t1, -# input_points.m1, -# input_points.m2, -# ] -# conditions = [input_points.pa, input_points.pb, input_points.p3] - -# p1, p2, det = mapping.map_forward(inputs, conditions) -# inv_inputs, inv_det = mapping.map_inverse([p1, p2], conditions) -# # r = inv_det * det -# # print("r stats:", np.min(r), np.median(r), np.max(r)) -# # print("abs(r-1) median:", np.median(np.abs(r-1))) -# # assert inv_det == approx(1 / det) -# for i, (inp, inv_inp) in enumerate(zip(inputs, inv_inputs)): -# if i == 0: -# assert ((inp < 0.5) == (inv_inp < 0.5)).all() -# continue -# assert inp == approx(inv_inp), f"mismatch in input index {i}" +def test_inverse(input_points): + mapping = ms.TwoToThreeParticleScattering() + + inputs = [ + input_points.r_choice, + input_points.r_s23, + input_points.r_t1, + input_points.m1, + input_points.m2, + ] + conditions = [input_points.pa, input_points.pb, input_points.p3] + + p1, p2, det = mapping.map_forward(inputs, conditions) + *inv_inputs, inv_det = mapping.map_inverse([p1, p2], conditions) + + assert inv_det == approx(1 / det, rel=1e-5) + + for i, (inp, inv_inp) in enumerate(zip(inputs, inv_inputs)): + if i == 0: + assert ((inp < 0.5) == (inv_inp < 0.5)).all() + continue + assert inp == approx(inv_inp), f"mismatch in input index {i}" def test_on_shell_masses(input_points): diff --git a/madspace/tests/test_color_ordered_mapping.py b/madspace/tests/test_color_ordered_mapping.py new file mode 100644 index 0000000000..953e7f5859 --- /dev/null +++ b/madspace/tests/test_color_ordered_mapping.py @@ -0,0 +1,209 @@ +import numpy as np +import pytest +from pytest import approx + +import madspace as ms + +# ---------------------------- +# Fixtures +# ---------------------------- + +@pytest.fixture +def rng(): + return np.random.default_rng(1234) + + +N = 10_000 # keep this moderate for CI speed; bump locally for tighter stats +CM_ENERGY = 13000.0 + + +# Mass set families. n_out is filled in at use. +def _massless(n_out): + return [0.0] * n_out, f"{n_out} massless" + + +def _w_like(n_out): + return [80.0] * n_out, f"{n_out} W-like" + + +def _mixed_top(n_out): + # one heavy, one W-like, rest massless; only meaningful for n_out >= 3. + return [173.0, 80.0] + [0.0] * (n_out - 2), f"{n_out} mixed" + + +def _mass_sets_for(n_out): + sets = [_massless(n_out), _w_like(n_out)] + if n_out >= 3: + sets.append(_mixed_top(n_out)) + return sets + + +# Each entry: (color_order, int_order_for_old_class, topology_label). +TOPOLOGIES = [ + ([0, 2, 3, 1, 4], [0, 1], "n=5: set1={2,3}, set2={4}"), + ([0, 2, 3, 4, 1, 5], [0, 1, 2], "n=6: set1={2,3,4}, set2={5}"), + ([0, 2, 3, 1, 4, 5], [0, 1, 2], "n=6: set1={2,3}, set2={4,5}"), + ([0, 2, 3, 4, 1, 5, 6], [0, 1, 2, 3], "n=7: set1={2,3,4}, set2={5,6}"), + ([0, 2, 3, 4, 5, 1, 6], [0, 1, 2, 3], "n=7: set1={2,3,4,5}, set2={6}"), + ([0, 2, 3, 4, 5, 6, 1, 7], [0, 1, 2, 3, 4], "n=8: set1={2,3,4,5,6}, set2={7}"), +] + + +def _expand(topologies): + """Flat list of (color_order, int_order, masses, label) pairs.""" + out = [] + for color_order, int_order, topo_label in topologies: + n_out = len(color_order) - 2 + for masses, mass_label in _mass_sets_for(n_out): + label = f"{mass_label}-{topo_label}" + out.append((color_order, int_order, masses, label)) + return out + + +CASES = _expand(TOPOLOGIES) +CASE_IDS = [c[3] for c in CASES] + + +def _physical_mask(det, p_ext): + """A physical event has finite det and finite momenta on every leg.""" + mask = np.isfinite(det) + for p in p_ext: + mask &= np.all(np.isfinite(p), axis=1) + return mask + + +def _filter_physical(det, p_ext): + """Keep physical events. Fail only if literally none survive.""" + mask = _physical_mask(det, p_ext) + if not mask.any(): + pytest.fail(f"No physical events out of {len(det)}") + return mask, [p[mask] for p in p_ext], det[mask] + +# ---------------------------- +# Tests +# ---------------------------- + +@pytest.mark.parametrize("color_order,_int_order,masses,_label", CASES, ids=CASE_IDS) +def test_momentum_conservation(rng, color_order, _int_order, masses, _label): + mapping = ms.ColorOrderedMapping(color_order) + r = rng.random((N, mapping.random_dim())) + e_cm = np.full(N, CM_ENERGY) + cond = [e_cm] + [np.full(N, m) for m in masses] + + *p_ext, det = mapping.map_forward(list(r.T), cond) + p_ext = list(p_ext) + _, p_ext, det = _filter_physical(det, p_ext) + + p_in = p_ext[0] + p_ext[1] + p_out_sum = sum(p_ext[i] for i in range(2, len(p_ext))) + assert p_in == approx(p_out_sum, abs=1e-3, rel=1e-6) + + +@pytest.mark.parametrize("color_order,_int_order,masses,_label", CASES, ids=CASE_IDS) +def test_on_shell_masses(rng, color_order, _int_order, masses, _label): + mapping = ms.ColorOrderedMapping(color_order) + r = rng.random((N, mapping.random_dim())) + e_cm = np.full(N, CM_ENERGY) + cond = [e_cm] + [np.full(N, m) for m in masses] + + *p_ext, det = mapping.map_forward(list(r.T), cond) + p_ext = list(p_ext) + _, p_ext, det = _filter_physical(det, p_ext) + + for i, m in enumerate(masses): + p = p_ext[i + 2] + m_check = np.sqrt(np.maximum(0, + p[:, 0] ** 2 - np.sum(p[:, 1:] ** 2, axis=1) + )) + assert m_check == approx(m, abs=1e-2, rel=1e-3), \ + f"particle {i} mass off" + + +@pytest.mark.parametrize("color_order,_int_order,masses,_label", CASES, ids=CASE_IDS) +def test_inverse(rng, color_order, _int_order, masses, _label): + """Forward o inverse = identity, to FP precision modulo a small tail. + + A handful of events at the very edge of the physical region produce a + forward det at the lower extreme of the distribution; for those, the + round-trip accumulates non-negligible FP error in atan2/sqrt/boost. We + drop the bottom 0.1% by |det| (genuine boundary events, near-zero + weight in any MC integral) and check the rest. + """ + n_events = N + mapping = ms.ColorOrderedMapping(color_order) + r = rng.random((n_events, mapping.random_dim())) + e_cm = np.full(n_events, CM_ENERGY) + cond = [e_cm] + [np.full(n_events, m) for m in masses] + + *p_ext, det = mapping.map_forward(list(r.T), cond) + p_ext = list(p_ext) + mask, p_ext_phys, det_phys = _filter_physical(det, p_ext) + cond_phys = [c[mask] for c in cond] + r_phys = r[mask] + + *r_inv, det_inv = mapping.map_inverse(p_ext_phys, cond_phys) + + floor = np.quantile(np.abs(det_phys), 0.001) + keep = np.isfinite(det_inv) & (np.abs(det_phys) > floor) + for ri in r_inv: + keep &= np.isfinite(ri) + + det_kept = det_phys[keep] + det_inv_kept = det_inv[keep] + # 0.1% tolerance absorbs accumulated FP error from chained + # boosts/rotations/sqrts in the deepest-walk topologies. Physics + # correctness is the job of test_phase_space_volume_matches_old. + assert det_inv_kept == approx(1 / det_kept, rel=1e-3) + + # Some randoms are discrete-branch choices (r_choice from 2->3 rungs); + # the round-trip recovers only the bin (< 0.5 vs >= 0.5), not the exact + # value. Bin-membership comparison handles those. + for i, (r_i, r_inv_i) in enumerate(zip(list(r_phys.T), r_inv)): + r_i_k = r_i[keep] + r_inv_i_k = r_inv_i[keep] + diff = np.abs(r_i_k - r_inv_i_k) + bin_ok = (r_i_k < 0.5) == (r_inv_i_k < 0.5) + ok = (diff < 1e-4) | bin_ok + assert np.all(ok), ( + f"random index {i}: {np.sum(~ok)} mismatches, " + f"max diff = {diff[~ok].max() if np.any(~ok) else 0:.3e}" + ) + + +@pytest.mark.parametrize("color_order,int_order,masses,_label", CASES, ids=CASE_IDS) +def test_phase_space_volume_matches_old(rng, color_order, int_order, masses, _label): + """The new class should integrate to the same phase-space volume as the + old TPropagatorMapping (matrix element = 1). + + This is *the* correctness test: both classes are MC estimators of the + same n-body integral, and their means must agree within a few sigma. + test_inverse is a precision check; this one is the physics check. + """ + n_events = N + new = ms.ColorOrderedMapping(color_order) + old = ms.TPropagatorMapping(int_order) + + r_new = rng.random((n_events, new.random_dim())) + r_old = rng.random((n_events, old.random_dim())) + e_cm = np.full(n_events, CM_ENERGY) + cond = [e_cm] + [np.full(n_events, m) for m in masses] + + *_, det_new = new.map_forward(list(r_new.T), cond) + *_, det_old = old.map_forward(list(r_old.T), cond) + + det_new = det_new[np.isfinite(det_new) & (det_new > 0)] + det_old = det_old[np.isfinite(det_old) & (det_old > 0)] + if not (len(det_new) and len(det_old)): + pytest.fail("no physical events in one of the estimators") + + mean_new, mean_old = np.mean(det_new), np.mean(det_old) + err_new = np.std(det_new) / np.sqrt(len(det_new)) + err_old = np.std(det_old) / np.sqrt(len(det_old)) + err = np.sqrt(err_new ** 2 + err_old ** 2) + + diff = abs(mean_new - mean_old) + assert diff < 5 * err, ( + f"phase-space volumes differ: new = {mean_new:.4e} +/- {err_new:.4e}, " + f"old = {mean_old:.4e} +/- {err_old:.4e}, " + f"diff = {diff:.4e} ({diff/err:.2f} sigma)" + ) \ No newline at end of file diff --git a/madspace/tests/test_double_t.py b/madspace/tests/test_double_t.py new file mode 100644 index 0000000000..b6ca9ce385 --- /dev/null +++ b/madspace/tests/test_double_t.py @@ -0,0 +1,411 @@ +from collections import namedtuple + +import numpy as np +import pytest +from pytest import approx + +import madspace as ms + +# ---------------------------- +# Utilities +# ---------------------------- + + +def inv_mass_sq(p): + return p[..., 0] ** 2 - np.sum(p[..., 1:] ** 2, axis=-1) + + +def mass(p): + return np.sqrt(np.maximum(inv_mass_sq(p), 0.0)) + + +# ---------------------------- +# Fixtures +# ---------------------------- + +N = 10_000 + +InputPoint = namedtuple( + "InputPoint", + [ + "r_phi", + "r_t1", + "r_t2", + "m1", + "mir_min", + "pa", + "pb", + ], +) + + +@pytest.fixture +def rng(): + return np.random.default_rng(1234) + + +@pytest.fixture( + params=[False, True], + ids=["LAB", "COM"], +) +def input_points(rng, request): + """Generate (pa, pb) and a set of forward inputs for the DoubleT mapping. + + Picks m1, mir_min so that the kinematic ranges of t1 and t2 are non-empty. + """ + com = request.param + + pz = rng.uniform(1000.0, 4000.0, N) + e1 = pz # massless beams in COM + e2 = pz + pa = np.stack([e1, np.zeros(N), np.zeros(N), +pz], axis=1) + pb = np.stack([e2, np.zeros(N), np.zeros(N), -pz], axis=1) + + if not com: + # boost to a random lab frame (the same recipe used in the 2->3 test) + pa3 = np.array([[0, 0, 4000.0]]) + rng.normal(0.0, 500.0, (N, 3)) + ea = np.sqrt(np.sum(pa3**2, axis=1)) # massless + pa = np.concatenate([ea[:, None], pa3], axis=1) + pb3 = np.array([[0, 0, -4000.0]]) + rng.normal(0.0, 500.0, (N, 3)) + eb = np.sqrt(np.sum(pb3**2, axis=1)) + pb = np.concatenate([eb[:, None], pb3], axis=1) + + # m1: the fixed single-particle mass + # mir_min: lower bound on recoil mass (e.g. sum of constituent masses). + # Pick mir_min well below the sqrt(s) so there\'s room for the recoil. + m1 = rng.uniform(1.0, 50.0, N) + mir_min = rng.uniform(50.0, 200.0, N) + + r_phi = rng.random(N) + r_t1 = rng.random(N) + r_t2 = rng.random(N) + + return InputPoint(r_phi, r_t1, r_t2, m1, mir_min, pa, pb) + + +@pytest.fixture( + params=[ + (np.full(N, 5.0), np.full(N, 80.0)), + (np.full(N, 50.0), np.full(N, 150.0)), + (np.full(N, 100.0), np.full(N, 250.0)), + ], + ids=[ + "light single, light recoil min", + "medium single, medium recoil min", + "heavy single, heavy recoil min", + ], +) +def fixed_input_points(rng, request): + """Fixed (m1, mir_min) used by the phase-space-volume test.""" + m1, mir_min = request.param + + pz = np.full(N, 3000.0) + e1 = pz + e2 = pz + pa = np.stack([e1, np.zeros(N), np.zeros(N), +pz], axis=1) + pb = np.stack([e2, np.zeros(N), np.zeros(N), -pz], axis=1) + + r_phi = rng.random(N) + r_t1 = rng.random(N) + r_t2 = rng.random(N) + + return InputPoint(r_phi, r_t1, r_t2, m1, mir_min, pa, pb) + + +# ---------------------------- +# Tests +# ---------------------------- + + +def test_momentum_conservation(input_points): + mapping = ms.DoubleT() + + inputs = [ + input_points.r_phi, + input_points.r_t1, + input_points.r_t2, + ] + conditions = [ + input_points.pa, + input_points.pb, + input_points.m1, + input_points.mir_min, + ] + + p1, p2, det = mapping.map_forward(inputs, conditions) + assert p1 + p2 == approx(input_points.pa + input_points.pb) + + +def test_on_shell_mass(input_points): + mapping = ms.DoubleT() + + inputs = [ + input_points.r_phi, + input_points.r_t1, + input_points.r_t2, + ] + conditions = [ + input_points.pa, + input_points.pb, + input_points.m1, + input_points.mir_min, + ] + + p1, p2, det = mapping.map_forward(inputs, conditions) + assert mass(p1) == approx(input_points.m1) + + +def test_recoil_mass_minimum(input_points): + """The recoil mass m_ir should be >= mir_min, since that bound goes + into the t1/t2 ranges.""" + mapping = ms.DoubleT() + + inputs = [ + input_points.r_phi, + input_points.r_t1, + input_points.r_t2, + ] + conditions = [ + input_points.pa, + input_points.pb, + input_points.m1, + input_points.mir_min, + ] + + p1, p2, det = mapping.map_forward(inputs, conditions) + m_ir = mass(p2) + # allow a small numerical tolerance below mir_min + assert np.all(m_ir >= input_points.mir_min - 1e-6 * input_points.mir_min) + + +def test_inverse(input_points): + mapping = ms.DoubleT() + + inputs = [ + input_points.r_phi, + input_points.r_t1, + input_points.r_t2, + ] + conditions = [ + input_points.pa, + input_points.pb, + input_points.m1, + input_points.mir_min, + ] + + p1, p2, det = mapping.map_forward(inputs, conditions) + *inv_inputs, inv_det = mapping.map_inverse([p1, p2], conditions) + + assert inv_det == approx(1 / det, rel=1e-5) + + for i, (inp, inv_inp) in enumerate(zip(inputs, inv_inputs)): + assert inp == approx(inv_inp), f"mismatch in input index {i}" + + +def test_phase_space_volume_analytic(fixed_input_points): + """The integral of DoubleT\'s det over (r_phi, r_t1, r_t2) should equal + the 2-body phase-space volume integrated over m_ir^2 in [mir_min^2, + (sqrt(s)-m1)^2]: + + V = int_{mir_min^2}^{(sqrt(s)-m1)^2} dm_ir^2 * (pi/(2s)) * sqrt(lambda(s, m1^2, mir^2)) + + This matches the 2->2 volume convention (no (2pi)^n factors). + """ + from scipy.integrate import quad + + mapping = ms.DoubleT() + inputs = [ + fixed_input_points.r_phi, + fixed_input_points.r_t1, + fixed_input_points.r_t2, + ] + conditions = [ + fixed_input_points.pa, + fixed_input_points.pb, + fixed_input_points.m1, + fixed_input_points.mir_min, + ] + p1, p2, det = mapping.map_forward(inputs, conditions) + + # All entries of fixed_input_points have the same s, m1, mir_min by construction. + pa = fixed_input_points.pa[0] + pb = fixed_input_points.pb[0] + s = (pa[0] + pb[0]) ** 2 - np.sum((pa[1:] + pb[1:]) ** 2) + m1_val = fixed_input_points.m1[0] + mir_min_val = fixed_input_points.mir_min[0] + + def lam(s, m1_sq, m2_sq): + return (s - m1_sq - m2_sq) ** 2 - 4 * m1_sq * m2_sq + + def integrand(mir_sq): + l = lam(s, m1_val ** 2, mir_sq) + if l <= 0: + return 0.0 + return np.pi / (2 * s) * np.sqrt(l) + + analytic_volume, _ = quad( + integrand, + mir_min_val ** 2, + (np.sqrt(s) - m1_val) ** 2, + ) + + mc_volume = np.mean(det) + mc_err = np.std(det) / np.sqrt(N) + + # 5-sigma agreement + diff = abs(mc_volume - analytic_volume) + assert diff < 5 * mc_err, ( + f"DoubleT volume mismatch: MC = {mc_volume:.6e} +/- {mc_err:.6e}, " + f"analytic = {analytic_volume:.6e}, " + f"diff = {diff:.4e} ({diff/mc_err:.2f} sigma), " + f"ratio MC/analytic = {mc_volume/analytic_volume:.4f}" + ) + + +def test_inverse_lab_specifically(rng): + """Explicit LAB test: round-trip should work in any frame, not just COM. + + This was the failing case for the 2->3 mapping before we fixed the + rotation; keep it as a guard for double_t too. + """ + mapping = ms.DoubleT() + + # boosted lab frame + pz = rng.uniform(1000.0, 4000.0, N) + pa3 = np.array([[0, 0, 4000.0]]) + rng.normal(0.0, 500.0, (N, 3)) + ea = np.sqrt(np.sum(pa3**2, axis=1)) + pa = np.concatenate([ea[:, None], pa3], axis=1) + pb3 = np.array([[0, 0, -4000.0]]) + rng.normal(0.0, 500.0, (N, 3)) + eb = np.sqrt(np.sum(pb3**2, axis=1)) + pb = np.concatenate([eb[:, None], pb3], axis=1) + + m1 = np.full(N, 30.0) + mir_min = np.full(N, 100.0) + r_phi = rng.random(N) + r_t1 = rng.random(N) + r_t2 = rng.random(N) + + inputs = [r_phi, r_t1, r_t2] + conditions = [pa, pb, m1, mir_min] + + p1, p2, det = mapping.map_forward(inputs, conditions) + *inv_inputs, inv_det = mapping.map_inverse([p1, p2], conditions) + + names = ["r_phi", "r_t1", "r_t2"] + for name, inp, inv_inp in zip(names, inputs, inv_inputs): + assert inp == approx(inv_inp), f"LAB mismatch in {name}" + assert inv_det == approx(1 / det, rel=1e-5) + + +# ---------------------------- +# Massless m1 stress tests (regression for the t-bound precision issue) +# ---------------------------- + + +@pytest.fixture( + params=[ + ("m1=0", 0.0), + ("m1=1e-9", 1e-9), + ("m1=1e-6", 1e-6), + ("m1=1e-3", 1e-3), + ], + ids=["m1=0", "m1=1e-9", "m1=1e-6", "m1=1e-3"], +) +def massless_input_points(rng, request): + """Near-massless m1 stresses the inverse: lsquare(p1) suffers + catastrophic cancellation when m1 is recovered from it. With m1 as a + condition (Option B fix), this should round-trip cleanly.""" + _, m1_val = request.param + + pz = np.full(N, 6500.0) # 13 TeV + pa = np.stack([pz, np.zeros(N), np.zeros(N), +pz], axis=1) + pb = np.stack([pz, np.zeros(N), np.zeros(N), -pz], axis=1) + + m1 = np.full(N, m1_val) + mir_min = np.full(N, 50.0) + + r_phi = rng.random(N) + r_t1 = rng.random(N) + r_t2 = rng.random(N) + + return InputPoint(r_phi, r_t1, r_t2, m1, mir_min, pa, pb) + + +def test_inverse_massless(massless_input_points): + """The DoubleT inverse must round-trip even for m1 ~= 0.""" + mapping = ms.DoubleT() + + inputs = [ + massless_input_points.r_phi, + massless_input_points.r_t1, + massless_input_points.r_t2, + ] + conditions = [ + massless_input_points.pa, + massless_input_points.pb, + massless_input_points.m1, + massless_input_points.mir_min, + ] + + p1, p2, det = mapping.map_forward(inputs, conditions) + *inv_inputs, inv_det = mapping.map_inverse([p1, p2], conditions) + + # det round-trip + assert inv_det == approx(1 / det, rel=1e-5) + + # random round-trip; r_t2 used to be the most sensitive (its bounds + # depend on t1_abs which depended on the lossy m1^2 recompute). + for i, (inp, inv_inp) in enumerate(zip(inputs, inv_inputs)): + assert inp == approx(inv_inp, rel=1e-5, abs=1e-8), \ + f"mismatch in input index {i}" + + +def test_phase_space_volume_massless(rng): + """Volume integral should still match analytic for massless m1.""" + from scipy.integrate import quad + + m1_val = 0.0 + mir_min_val = 80.0 + s_val = 13000.0 ** 2 + + pa = np.tile([6500.0, 0, 0, 6500.0], (N, 1)) + pb = np.tile([6500.0, 0, 0, -6500.0], (N, 1)) + + m1 = np.full(N, m1_val) + mir_min = np.full(N, mir_min_val) + + r_phi = rng.random(N) + r_t1 = rng.random(N) + r_t2 = rng.random(N) + + mapping = ms.DoubleT() + p1, p2, det = mapping.map_forward( + [r_phi, r_t1, r_t2], + [pa, pb, m1, mir_min], + ) + + def lam(s, m1_sq, m2_sq): + return (s - m1_sq - m2_sq) ** 2 - 4 * m1_sq * m2_sq + + def integrand(mir_sq): + l = lam(s_val, m1_val ** 2, mir_sq) + if l <= 0: + return 0.0 + return np.pi / (2 * s_val) * np.sqrt(l) + + analytic_volume, _ = quad( + integrand, + mir_min_val ** 2, + (np.sqrt(s_val) - m1_val) ** 2, + ) + + mc_volume = np.mean(det) + mc_err = np.std(det) / np.sqrt(N) + + diff = abs(mc_volume - analytic_volume) + assert diff < 5 * mc_err, ( + f"DoubleT massless volume mismatch: MC = {mc_volume:.6e} +/- " + f"{mc_err:.6e}, analytic = {analytic_volume:.6e}, " + f"diff = {diff:.4e} ({diff/mc_err:.2f} sigma), " + f"ratio MC/analytic = {mc_volume/analytic_volume:.4f}" + ) \ No newline at end of file