This commit is contained in:
@@ -0,0 +1,451 @@
|
||||
# -*- coding: utf-8 -*- #
|
||||
# frozen_string_literal: true
|
||||
|
||||
module Rouge
|
||||
module Lexers
|
||||
class Stan < RegexLexer
|
||||
title "Stan"
|
||||
desc 'Stan Modeling Language (mc-stan.org)'
|
||||
tag 'stan'
|
||||
filenames '*.stan', '*.stanfunctions'
|
||||
|
||||
# optional comment or whitespace
|
||||
WS = %r((?:\s|//.*?\n|/[*].*?[*]/)+)
|
||||
ID = /[a-zA-Z_][a-zA-Z0-9_]*/
|
||||
RT = /(?:(?:[a-z_]\s*(?:\[[0-9, ]\])?)\s+)*/
|
||||
OP = Regexp.new([
|
||||
# Assigment operators
|
||||
"=",
|
||||
|
||||
# Comparison operators
|
||||
"<", "<=", ">", ">=", "==", "!=",
|
||||
|
||||
# Boolean operators
|
||||
"!", "&&", "\\|\\|",
|
||||
|
||||
# Real-valued arithmetic operators
|
||||
"\\+", "-", "\\*", "/", "\\^",
|
||||
|
||||
# Transposition operator
|
||||
"'",
|
||||
|
||||
# Elementwise functions
|
||||
"\\.\\+", "\\.-", "\\.\\*", "\\./", "\\.\\^",
|
||||
|
||||
# Matrix division operators
|
||||
"\\\\",
|
||||
|
||||
# Compound assigment operators
|
||||
"\\+=", "-=", "\\*=", "/=", "\\.\\*=", "\\./=",
|
||||
|
||||
# Sampling
|
||||
"~",
|
||||
|
||||
# Conditional operator
|
||||
"\\?", ":"
|
||||
].join("|"))
|
||||
|
||||
def self.keywords
|
||||
@keywords ||= Set.new %w(
|
||||
if else while for break continue print reject return
|
||||
)
|
||||
end
|
||||
|
||||
def self.types
|
||||
@types ||= Set.new %w(
|
||||
int real vector ordered positive_ordered simplex unit_vector
|
||||
row_vector matrix cholesky_factor_corr cholesky_factor_cov corr_matrix
|
||||
cov_matrix data void complex array
|
||||
)
|
||||
end
|
||||
|
||||
def self.reserved
|
||||
@reserved ||= Set.new [
|
||||
# Reserved words from Stan language
|
||||
"for", "in", "while", "repeat", "until", "if", "then", "else", "true",
|
||||
"false", "target", "functions", "model", "data", "parameters",
|
||||
"quantities", "transformed", "generated",
|
||||
|
||||
# Reserved names from Stan implementation
|
||||
"var", "fvar", "STAN_MAJOR", "STAN_MINOR", "STAN_PATCH",
|
||||
"STAN_MATH_MAJOR", "STAN_MATH_MINOR", "STAN_MATH_PATCH",
|
||||
|
||||
# Reserved names from C++
|
||||
"alignas", "alignof", "and", "and_eq", "asm", "auto", "bitand",
|
||||
"bitor", "bool", "break", "case", "catch", "char", "char16_t",
|
||||
"char32_t", "class", "compl", "const", "constexpr", "const_cast",
|
||||
"continue", "decltype", "default", "delete", "do", "double",
|
||||
"dynamic_cast", "else", "enum", "explicit", "export", "extern",
|
||||
"false", "float", "for", "friend", "goto", "if", "inline", "int",
|
||||
"long", "mutable", "namespace", "new", "noexcept", "not", "not_eq",
|
||||
"nullptr", "operator", "or", "or_eq", "private", "protected",
|
||||
"public", "register", "reinterpret_cast", "return", "short", "signed",
|
||||
"sizeof", "static", "static_assert", "static_cast", "struct",
|
||||
"switch", "template", "this", "thread_local", "throw", "true", "try",
|
||||
"typedef", "typeid", "typename", "union", "unsigned", "using",
|
||||
"virtual", "void", "volatile", "wchar_t", "while", "xor", "xor_eq"
|
||||
]
|
||||
end
|
||||
|
||||
def self.builtin_functions
|
||||
@builtin_functions ||= Set.new [
|
||||
# Integer-Valued Basic Functions
|
||||
|
||||
## Absolute functions
|
||||
"abs", "int_step",
|
||||
|
||||
## Bound functions
|
||||
"min", "max",
|
||||
|
||||
## Size functions
|
||||
"size",
|
||||
|
||||
# Real-Valued Basic Functions
|
||||
|
||||
## Log probability function
|
||||
"target", "get_lp",
|
||||
|
||||
## Logical functions
|
||||
"step", "is_inf", "is_nan",
|
||||
|
||||
## Step-like functions
|
||||
"fabs", "fdim", "fmin", "fmax", "fmod", "floor", "ceil", "round",
|
||||
"trunc",
|
||||
|
||||
## Power and logarithm functions
|
||||
"sqrt", "cbrt", "square", "exp", "exp2", "log", "log2", "log10",
|
||||
"pow", "inv", "inv_sqrt", "inv_square",
|
||||
|
||||
## Trigonometric functions
|
||||
"hypot", "cos", "sin", "tan", "acos", "asin", "atan", "atan2",
|
||||
|
||||
## Hyperbolic trigonometric functions
|
||||
"cosh", "sinh", "tanh", "acosh", "asinh", "atanh",
|
||||
|
||||
## Link functions
|
||||
"logit", "inv_logit", "inv_cloglog",
|
||||
|
||||
## Probability-related functions
|
||||
"erf", "erfc", "Phi", "inv_Phi", "Phi_approx", "binary_log_loss",
|
||||
"owens_t",
|
||||
|
||||
## Combinatorial functions
|
||||
"beta", "inc_beta", "lbeta", "tgamma", "lgamma", "digamma",
|
||||
"trigamma", "lmgamma", "gamma_p", "gamma_q",
|
||||
"binomial_coefficient_log", "choose", "bessel_first_kind",
|
||||
"bessel_second_kind", "modified_bessel_first_kind",
|
||||
"log_modified_bessel_first_kind", "modified_bessel_second_kind",
|
||||
"falling_factorial", "lchoose", "log_falling_factorial",
|
||||
"rising_factorial", "log_rising_factorial",
|
||||
|
||||
## Composed functions
|
||||
"expm1", "fma", "multiply_log", "ldexp", "lmultiply", "log1p",
|
||||
"log1m", "log1p_exp", "log1m_exp", "log_diff_exp", "log_mix",
|
||||
"log_sum_exp", "log_inv_logit", "log_inv_logit_diff",
|
||||
"log1m_inv_logit",
|
||||
|
||||
## Special functions
|
||||
"lambert_w0", "lambert_wm1",
|
||||
|
||||
# Complex-Valued Basic Functions
|
||||
|
||||
## Complex constructors and accessors
|
||||
"to_complex", "get_real", "get_imag",
|
||||
|
||||
## Complex special functions
|
||||
"arg", "norm", "conj", "proj", "polar",
|
||||
|
||||
# Array Operations
|
||||
|
||||
## Reductions
|
||||
"sum", "prod", "log_sum_exp", "mean", "variance", "sd", "distance",
|
||||
"squared_distance", "quantile",
|
||||
|
||||
## Array size and dimension function
|
||||
"dims", "num_elements",
|
||||
|
||||
## Array broadcasting
|
||||
"rep_array",
|
||||
|
||||
## Array concatenation
|
||||
"append_array",
|
||||
|
||||
## Sorting functions
|
||||
"sort_asc", "sort_desc", "sort_indices_asc", "sort_indices_desc",
|
||||
"rank",
|
||||
|
||||
## Reversing functions
|
||||
"reverse",
|
||||
|
||||
# Matrix Operations
|
||||
|
||||
## Integer-valued matrix size functions
|
||||
"num_elements", "rows", "cols",
|
||||
|
||||
## Dot products and specialized products
|
||||
"dot_product", "columns_dot_product", "rows_dot_product", "dot_self",
|
||||
"columns_dot_self", "rows_dot_self", "tcrossprod", "crossprod",
|
||||
"quad_form", "quad_form_diag", "quad_form_sym", "trace_quad_form",
|
||||
"trace_gen_quad_form", "multiply_lower_tri_self_transpose",
|
||||
"diag_pre_multiply", "diag_post_multiply",
|
||||
|
||||
## Broadcast functions
|
||||
"rep_vector", "rep_row_vector", "rep_matrix",
|
||||
"symmetrize_from_lower_tri",
|
||||
|
||||
## Diagonal matrix functions
|
||||
"add_diag", "diagonal", "diag_matrix", "identity_matrix",
|
||||
|
||||
## Container construction functions
|
||||
"linspaced_array", "linspaced_int_array", "linspaced_vector",
|
||||
"linspaced_row_vector", "one_hot_int_array", "one_hot_array",
|
||||
"one_hot_vector", "one_hot_row_vector", "ones_int_array",
|
||||
"ones_array", "ones_vector", "ones_row_vector", "zeros_int_array",
|
||||
"zeros_array", "zeros_vector", "zeros_row_vector", "uniform_simplex",
|
||||
|
||||
## Slicing and blocking functions
|
||||
"col", "row", "block", "sub_col", "sub_row", "head", "tail",
|
||||
"segment",
|
||||
|
||||
## Matrix concatenation
|
||||
"append_col", "append_row",
|
||||
|
||||
## Special matrix functions
|
||||
"softmax", "log_softmax", "cumulative_sum",
|
||||
|
||||
## Covariance functions
|
||||
"cov_exp_quad",
|
||||
|
||||
## Linear algebra functions and solvers
|
||||
"mdivide_left_tri_low", "mdivide_right_tri_low", "mdivide_left_spd",
|
||||
"mdivide_right_spd", "matrix_exp", "matrix_exp_multiply",
|
||||
"scale_matrix_exp_multiply", "matrix_power", "trace", "determinant",
|
||||
"log_determinant", "inverse", "inverse_spd", "chol2inv",
|
||||
"generalized_inverse", "eigenvalues_sym", "eigenvectors_sym",
|
||||
"qr_thin_Q", "qr_thin_R", "qr_Q", "qr_R", "cholseky_decompose",
|
||||
"singular_values", "svd_U", "svd_V",
|
||||
|
||||
# Sparse Matrix Operations
|
||||
|
||||
## Conversion functions
|
||||
"csr_extract_w", "csr_extract_v", "csr_extract_u",
|
||||
"csr_to_dense_matrix",
|
||||
|
||||
## Sparse matrix arithmetic
|
||||
"csr_matrix_times_vector",
|
||||
|
||||
# Mixed Operations
|
||||
"to_matrix", "to_vector", "to_row_vector", "to_array_2d",
|
||||
"to_array_1d",
|
||||
|
||||
# Higher-Order Functions
|
||||
|
||||
## Algebraic equation solver
|
||||
"algebra_solver", "algebra_solver_newton",
|
||||
|
||||
## Ordinary differential equation
|
||||
"ode_rk45", "ode_rk45_tol", "ode_ckrk", "ode_ckrk_tol", "ode_adams",
|
||||
"ode_adams_tol", "ode_bdf", "ode_bdf_tol", "ode_adjoint_tol_ctl",
|
||||
|
||||
## 1D integrator
|
||||
"integrate_1d",
|
||||
|
||||
## Reduce-sum function
|
||||
"reduce_sum", "reduce_sum_static",
|
||||
|
||||
## Map-rect function
|
||||
"map_rect",
|
||||
|
||||
# Deprecated Functions
|
||||
"integrate_ode_rk45", "integrate_ode", "integrate_ode_adams",
|
||||
"integrate_ode_bdf",
|
||||
|
||||
# Hidden Markov Models
|
||||
"hmm_marginal", "hmm_latent_rng", "hmm_hidden_state_prob"
|
||||
]
|
||||
end
|
||||
|
||||
def self.distributions
|
||||
@distributions ||= Set.new(
|
||||
[
|
||||
# Discrete Distributions
|
||||
|
||||
## Binary Distributions
|
||||
"bernoulli", "bernoulli_logit", "bernoulli_logit_glm",
|
||||
|
||||
## Bounded Discrete Distributions
|
||||
"binomial", "binomial_logit", "beta_binomial", "hypergeometric",
|
||||
"categorical", "categorical_logit_glm", "discrete_range",
|
||||
"ordered_logistic", "ordered_logistic_glm", "ordered_probit",
|
||||
|
||||
## Unbounded Discrete Distributions
|
||||
"neg_binomial", "neg_binomial_2", "neg_binomial_2_log",
|
||||
"neg_binomial_2_log_glm", "poisson", "poisson_log",
|
||||
"poisson_log_glm",
|
||||
|
||||
## Multivariate Discrete Distributions
|
||||
"multinomial", "multinomial_logit",
|
||||
|
||||
# Continuous Distributions
|
||||
|
||||
## Unbounded Continuous Distributions
|
||||
"normal", "std_normal", "normal_id_glm", "exp_mod_normal",
|
||||
"skew_normal", "student_t", "cauchy", "double_exponential",
|
||||
"logistic", "gumbel", "skew_double_exponential",
|
||||
|
||||
## Positive Continuous Distributions
|
||||
"lognormal", "chi_square", "inv_chi_square",
|
||||
"scaled_inv_chi_square", "exponential", "gamma", "inv_gamma",
|
||||
"weibull", "frechet", "rayleigh",
|
||||
|
||||
## Positive Lower-Bounded Distributions
|
||||
"pareto", "pareto_type_2", "wiener",
|
||||
|
||||
## Continuous Distributions on [0, 1]
|
||||
"beta", "beta_proportion",
|
||||
|
||||
## Circular Distributions
|
||||
"von_mises",
|
||||
|
||||
## Bounded Continuous Distributions
|
||||
"uniform",
|
||||
|
||||
## Distributions over Unbounded Vectors
|
||||
"multi_normal", "multi_normal_prec", "multi_normal_cholesky",
|
||||
"multi_gp", "multi_gp_cholesky", "multi_student_t",
|
||||
"gaussian_dlm_obs",
|
||||
|
||||
## Simplex Distributions
|
||||
"dirichlet",
|
||||
|
||||
## Correlation Matrix Distributions
|
||||
"lkj_corr", "lkj_corr_cholesky",
|
||||
|
||||
## Covariance Matrix Distributions
|
||||
"wishart", "inv_wishart"
|
||||
].product([
|
||||
"", "_lpmf", "_lupmf", "_lpdf", "_lcdf", "_lccdf", "_rng", "_log",
|
||||
"_cdf_log", "_ccdf_log"
|
||||
]).map {|s| "#{s[0]}#{s[1]}"}
|
||||
)
|
||||
end
|
||||
|
||||
def self.constants
|
||||
@constants ||= Set.new [
|
||||
# Mathematical constants
|
||||
"pi", "e", "sqrt2", "log2", "log10",
|
||||
|
||||
# Special values
|
||||
"not_a_number", "positive_infinity", "negative_infinity",
|
||||
"machine_precision"
|
||||
]
|
||||
end
|
||||
|
||||
state :root do
|
||||
mixin :whitespace
|
||||
rule %r/#include/, Comment::Preproc, :include
|
||||
rule %r/#.*$/, Generic::Deleted
|
||||
rule %r(
|
||||
functions
|
||||
|(?:transformed\s+)?data
|
||||
|(?:transformed\s+)?parameters
|
||||
|model
|
||||
|generated\s+quantities
|
||||
)x, Name::Namespace
|
||||
rule %r(\{), Punctuation, :bracket_scope
|
||||
mixin :scope
|
||||
end
|
||||
|
||||
state :include do
|
||||
rule %r((\s+)(\S+)(\s*)) do |m|
|
||||
token Text, m[1]
|
||||
token Comment::PreprocFile, m[2]
|
||||
token Text, m[3]
|
||||
pop!
|
||||
end
|
||||
end
|
||||
|
||||
state :whitespace do
|
||||
rule %r(\n+)m, Text
|
||||
rule %r(//(\\.|.)*?$), Comment::Single
|
||||
mixin :inline_whitespace
|
||||
end
|
||||
|
||||
state :inline_whitespace do
|
||||
rule %r([ \t\r]+), Text
|
||||
rule %r(/(\\\n)?[*].*?[*](\\\n)?/)m, Comment::Multiline
|
||||
end
|
||||
|
||||
state :statements do
|
||||
mixin :whitespace
|
||||
rule %r/#include/, Comment::Preproc, :include
|
||||
rule %r/#.*$/, Generic::Deleted
|
||||
rule %r("), Str, :string
|
||||
rule %r(
|
||||
(
|
||||
((\d+[.]\d*|[.]?\d+)e[+-]?\d+|\d*[.]\d+|\d+)
|
||||
(#{WS})[+-](#{WS})
|
||||
((\d+[.]\d*|[.]?\d+)e[+-]?\d+|\d*[.]\d+|\d+)i
|
||||
)
|
||||
|((\d+[.]\d*|[.]?\d+)e[+-]?\d+|\d*[.]\d+|\d+)i
|
||||
|((\d+[.]\d*|[.]?\d+)e[+-]?\d+|\d*[.]\d+)
|
||||
)mx, Num::Float
|
||||
rule %r/\d+/, Num::Integer
|
||||
rule %r(\*/), Error
|
||||
rule OP, Operator
|
||||
rule %r([\[\],.;]), Punctuation
|
||||
rule %r([|](?![|])), Punctuation
|
||||
rule %r(T\b), Keyword::Reserved
|
||||
rule %r((lower|upper)\b), Name::Attribute
|
||||
rule ID do |m|
|
||||
name = m[0]
|
||||
|
||||
if self.class.keywords.include? name
|
||||
token Keyword
|
||||
elsif self.class.types.include? name
|
||||
token Keyword::Type
|
||||
elsif self.class.reserved.include? name
|
||||
token Keyword::Reserved
|
||||
else
|
||||
token Name::Variable
|
||||
end
|
||||
end
|
||||
end
|
||||
|
||||
state :scope do
|
||||
mixin :whitespace
|
||||
rule %r(
|
||||
(#{RT}) # Return type
|
||||
(#{ID}) # Function name
|
||||
(?=\([^;]*?\)) # Signature or arguments
|
||||
)mx do |m|
|
||||
recurse m[1]
|
||||
|
||||
name = m[2]
|
||||
if self.class.builtin_functions.include? name
|
||||
token Name::Builtin, name
|
||||
elsif self.class.distributions.include? name
|
||||
token Name::Builtin, name
|
||||
elsif self.class.constants.include? name
|
||||
token Keyword::Constant
|
||||
else
|
||||
token Name::Function, name
|
||||
end
|
||||
end
|
||||
rule %r(\{), Punctuation, :bracket_scope
|
||||
rule %r(\(), Punctuation, :parens_scope
|
||||
mixin :statements
|
||||
end
|
||||
|
||||
state :bracket_scope do
|
||||
mixin :scope
|
||||
rule %r(\}), Punctuation, :pop!
|
||||
end
|
||||
|
||||
state :parens_scope do
|
||||
mixin :scope
|
||||
rule %r(\)), Punctuation, :pop!
|
||||
end
|
||||
end
|
||||
end
|
||||
end
|
||||
Reference in New Issue
Block a user