R/util.affinity.R

In CHNOSZ: Thermodynamic Calculations and Diagrams for Geochemistry

# CHNOSZ/util-affinity.R
# Helper functions for affinity()

## If this file is interactively sourced, the following are also needed to provide unexported functions:
#source("util.units.R")

slice.affinity <- function(affinity,d=1,i=1) {
  # Take a slice of affinity along one dimension
  a <- affinity
  for(j in 1:length(a$values)) {
    # Preserve the dimensions (especially: names(mydim))
    # - fix for change in behavior of aperm in R-devel 2015-11-17
    mydim <- dim(a$values[[j]])
    a$values[[j]] <- as.array(slice(a$values[[j]],d=d,i=i))
    # The dimension from which we take the slice vanishes
    dim(a$values[[j]]) <- mydim[-d]
  }
  return(a)
}

### Unexported functions ###

energy <- function(what,vars,vals,lims,T=298.15,P="Psat",IS=0,sout=NULL,exceed.Ttr=FALSE,exceed.rhomin=FALSE,transect=FALSE) {
  # 20090329 Dxtracted from affinity() and made to
  # deal with >2 dimensions (variables)

  # Calculate "what" property
  # logK - logK of formation reactions
  # logact.basis - logarithms of activities of basis species in reactions
  # logact.species - logarithms of activities of species in reactions
  # logQ - logQ of formation reactions
  # A - A/2.303RT of formation reactions
  # "vars": vector of variables (T,P,basisnames,IS)
  # "vals": list of values for each variable
  # "lims": used to get the dimensions (resolution for each variable)

  ### Some argument checking
  if(length(unique(vars)) != length(vars)) stop("please supply unique variable names")
  ## Basis definition / number of basis species 
  mybasis <- basis()
  nbasis <- nrow(mybasis)
  ## Species definition / number of species
  myspecies <- get("thermo", CHNOSZ)$species
  if(is.character(what)) {
    if(is.null(myspecies)) stop('species properties requested, but species have not been defined')
    nspecies <- nrow(myspecies)
    if(!identical(what,"logact.basis")) ispecies <- 1:nspecies
  }
  ## The dimensions of our arrays
  resfun <- function(lim) lim[3]
  mydim <- sapply(lims,resfun)
  if(transect) {
    if(!isTRUE(all.equal(min(mydim),max(mydim)))) stop("variables define a transect but their lengths are not all equal")
    mydim <- mydim[[1]]
  }
  ## The number of dimensions we have
  if(transect) nd <- 1 else nd <- length(vars) # == length(mydim)
  ## Basis names / which vars denote basis species
  basisnames <- rownames(mybasis)
  ibasisvar <- match(vars,basisnames)
  varisbasis <- !is.na(ibasisvar)
  ibasisvar <- ibasisvar[!is.na(ibasisvar)]
  ## Which vars are in P,T,IS
  varissubcrt <- vars %in% c("P","T","IS")
  if(sum(varissubcrt) > 2) stop("only up to 2 of P,T,IS are supported")
  ## Categorize the basis species:
  # 0 - not in the vars; 1 - one of the vars
  ibasis <- 1:nbasis
  ibasis0 <- ibasis[!ibasis %in% ibasisvar]
  ibasis1 <- ibasis[ibasis %in% ibasisvar]
  if(identical(what,"logact.basis")) ispecies <- ibasis
  ## What subcrt variable is used to make a 2-D grid?
  workaround.IS <- FALSE
  if(sum(varissubcrt) > 1 & !transect) {
    grid <- vars[varissubcrt][1]
    if("IS" %in% vars) {
      if(grid != "IS") workaround.IS <- TRUE
      grid <- "IS"
    }
  } else grid <- NULL
  ### Done argument processing

  ### Function to index the variables in a permuted order
  # by swapping ivar for the first
  # e.g. for nd=5, ivars(4)=c(4,2,3,1,5)
  ivars <- function(ivar,iv=NULL) {
    if(nd==0) return(1)
    if(is.null(iv)) iv <- 1:nd
    iv.1 <- iv[1]
    iv[1] <- ivar
    iv[ivar] <- iv.1
    return(iv)
  }

  ### Functions for logact / logQ
  # A basis species not in var
  logact.basis0.fun <- function(ibasis) {
    logact <- mybasis$logact[ibasis]
    # for the case where this basis species is buffered
    if(!can.be.numeric(logact)) logact <- 0
    else logact <- as.numeric(logact)
    return(array(logact,mydim))
  }
  # A basis species in var
  logact.basis1.fun <- function(ivar) {
    dim.in <- dim(vals[[ivar]])
    if(is.null(dim.in)) dim.in <- 0
    # why doesn't this work?
    # if(identical(dim.in,mydim))
    if(all(dim.in==mydim) | transect) return(vals[[ivar]])
    else return(aperm(array(vals[[ivar]],mydim[ivars(ivar)]),ivars(ivar)))
  }
  # Any basis species
  logact.basis.fun <- function(ibasis) {
    if(ibasis %in% ibasis0) return(logact.basis0.fun(ibasis))
    else return(logact.basis1.fun(match(basisnames[ibasis],vars)))
  }
  # All basis species
  logact.basis <- function() lapply(ibasis,logact.basis.fun)
  # logact of a single species
  logact.species.fun <- function(ispecies) array(myspecies$logact[ispecies],mydim)
  ## Contributions to logQ
  # From a single basis species 
  logQ.basis.fun <- function(ibasis,coeff) - coeff * logact.basis.fun(ibasis)
  # From all basis species in a single formation reaction
  logQ.basis.species <- function(ispecies) 
    Reduce("+", mapply(logQ.basis.fun,ibasis,myspecies[ispecies,1:nbasis],SIMPLIFY=FALSE))
  # All basis species in all reactions
  logQ.basis <- function() mapply(logQ.basis.species,1:nspecies,SIMPLIFY=FALSE)
  # By a single species
  logQ.species.fun <- function(ispecies,coeff) coeff * logact.species.fun(ispecies)
  # By all species
  logQ.species <- function() 
    mapply(logQ.species.fun,1:nspecies,1,SIMPLIFY=FALSE)
  # Total logQ of all reactions
  logQ <- function() lsum(logQ.basis(),logQ.species())

  ### Function for calling subcrt
  sout.fun <- function(property="logK") {
    species <- c(mybasis$ispecies,myspecies$ispecies)
    if(!is.null(sout)) {
      # Extract the needed species from a provided sout 20190131
      isout <- match(species, sout$species$ispecies)
      this.sout <- sout
      this.sout$species <- this.sout$species[isout, ]
      this.sout$out <- this.sout$out[isout]
      return(this.sout) 
    } else {
      ## subcrt arguments
      if("T" %in% vars) T <- vals[[which(vars=="T")]]
      if("P" %in% vars) P <- vals[[which(vars=="P")]]
      if("IS" %in% vars) IS <- vals[[which(vars=="IS")]]
      s.args <- list(species=species,property=property,T=T,P=P,IS=IS,grid=grid,convert=FALSE,exceed.Ttr=exceed.Ttr,exceed.rhomin=exceed.rhomin)
      sout <- do.call("subcrt",s.args)
      # 20191210 workaround a limitation in subcrt(): only IS (if present) can be the 'grid' variable
      if(workaround.IS) {
        lenIS <- length(IS)
        # Only T or P has length > 1
        lenTP <- max(length(T), length(P))
        # Make an indexing matrix "byrow" then vectorize it (not byrow) to get permutation indices
        indmat <- matrix(1:(lenIS * lenTP), nrow = lenIS, byrow = TRUE)
        ind <- as.vector(indmat)
        sout$out <- lapply(sout$out, function(x) x[indmat, ])
      }
      # Species indices are updated by subcrt() for minerals with phase transitions
      # e.g. i <- info("chalcocite"); subcrt(i, T=200)$species$ispecies == i + 1
      # so we should keep the original species index to be able to find the species in a provided 'sout'
      # (noted for Mosaic diagram section of anintro.Rmd 20190203)
      sout$species$ispecies <- species
      return(sout)
    }
  }

  ### Functions for logK/subcrt props
  # The logK contribution by any species or basis species
  X.species <- function(ispecies,coeff,X) coeff * sout$out[[ispecies]][,names(sout$out[[ispecies]])==X]
  # The logK contribution by all basis species in a reaction
  X.basis <- function(ispecies,X) Reduce("+", mapply(X.species,ibasis,-myspecies[ispecies,ibasis],X,SIMPLIFY=FALSE))
  # The logK of any reaction
  X.reaction <- function(ispecies,X) X.species((ispecies+nbasis),1,X) + X.basis(ispecies,X)
  # To get logK or subcrt props or other values into the dimensions we are using
  dim.fun <- function(x,idim=NULL) {
    if(is.null(idim)) {
      if(transect) idim <- 1
      else if(is.null(grid)) {
        # One of T,P,IS
        ivar <- which(vars %in% c("T","P","IS"))
        if(length(ivar)==0) ivar <- 1
        idim <- ivars(ivar)
      } else {
        # Two of T,P,IS
        ivar1 <- which(varissubcrt)[1]
        ivar2 <- which(varissubcrt)[2]
        idim <- ivars(ivar2, ivars(ivar1))
      }
    }
    return(aperm(array(x,mydim[idim]),idim))
  }
  # Properties of all species
  X.fun <- function(X) lapply(lapply(ispecies,X.reaction,X),dim.fun)
  logK <- function() lapply(ispecies,X.reaction,"logK")
  # A/2.303RT
  A <- function() {
    out <- mapply(`-`, X.fun("logK"), logQ(), SIMPLIFY = FALSE)
    # Deal with affinities of protein ionization here 20120527
    if("H+" %in% rownames(mybasis)) {
      # Which species are proteins
      isprotein <- grepl("_", myspecies$name)
      if(any(isprotein)) {
        if(get("thermo", CHNOSZ)$opt$ionize.aa) {
            message("affinity: ionizing proteins ...")
            # The rownumbers in thermo()$protein
            ip <- pinfo(myspecies$name[isprotein])
            # Get the affinity of ionization
            iHplus <- match("H+", rownames(mybasis))
            # as.numeric is needed in case the logact column is character mode
            # due to buffer definition (that means we can't do pH buffers)
            pH <- -as.numeric(mybasis$logact[iHplus])
            A.ionization <- A.ionization(ip, vars, vals, T = T, P = P, pH = pH, transect = transect)
            # Add it to the affinities of formation reactions of the non-ionized proteins
            out[isprotein] <- lsum(out[isprotein], A.ionization)
        } else message("affinity: NOT ionizing proteins because thermo()$opt$ionize.aa is FALSE")
      }
    }
    return(out)
  }

  ### Call the necessary functions
  if(!is.character(what)) {
    # Expand numeric values into our dimensions
    # (used by energy.args() for calculating pe=f(Eh,T) )
    # TODO: document that sout here denotes the dimension
    # we're expanding into
    return(dim.fun(what,ivars(sout$out)))
  } else if(what %in% c('G','H','S','Cp','V','E','kT','logK')) {
    # Get subcrt properties for reactions
    sout <- sout.fun(what)
    a <- X.fun(what)
  } else if(length(agrep("species",what)) > 0) {
    # Get subcrt properties for species
    # e.g. what=G.species, Cp.species etc
    mywhat <- s2c(what,sep=".",keep.sep=FALSE)[1]
    sout <- sout.fun(mywhat)
    a <- lapply(mapply(X.species,ispecies+nbasis,1,mywhat,SIMPLIFY=FALSE),dim.fun)
  } else {
    # Get affinities or logact.basis, logact.species or logQ
    if(what=="A") sout <- sout.fun("logK")
    a <- do.call(what,list())
  }

  ### Use species indices as the names 
  if(what=="logact.basis") names(a) <- mybasis$ispecies
  else names(a) <- myspecies$ispecies
  ### Return the results
  return(list(sout=sout,a=a))
}

energy.args <- function(args, transect = FALSE) {
  ## Extracted from affinity() and modified 20090329 jmd
  # Takes a list of arguments which define the dimensions
  # over which to calculate logQ, logK and affinity
  # The names should be T, P, IS and names of basis species
  # (or pH, pe, Eh)
  thermo <- get("thermo", CHNOSZ)
  ## Inputs are like c(T1,T2,res)
  # and outputs are like seq(T1,T2,length.out=res)
  # unless transect: do the variables specify a transect? 20090627
  transect <- any(transect, any(sapply(args,length) > 3))
  ## Convert T, P args and take care of 
  # Single values for T, P or IS
  T <- 298.15
  P <- "Psat"
  IS <- 0
  T.is.var <- P.is.var <- IS.is.var <- FALSE
  arg.is.T <- names(args)=="T"
  arg.is.P <- names(args)=="P"
  arg.is.IS <- names(args)=="IS"
  if(any(arg.is.T)) {
    T <- args[[which(arg.is.T)]]
    if(length(T) > 1) T.is.var <- TRUE
    if(transect) args[[which(arg.is.T)]] <- T <- envert(T,'K')
    else args[[which(arg.is.T)]][1:2] <- T[1:2] <- envert(T[1:2],'K')
    T <- T[!is.na(T)]
  }
  if(any(arg.is.P)) {
    P <- args[[which(arg.is.P)]]
    if(length(P) > 1) P.is.var <- TRUE
    if(!identical(P,"Psat")) {
      if(transect) args[[which(arg.is.P)]] <- P <- envert(P,'bar')
      else args[[which(arg.is.P)]][1:2] <- P[1:2] <- envert(P[1:2],'bar')
    }
    P <- P[!is.na(P)]
  }
  if(any(arg.is.IS)) {
    IS <- args[[which(arg.is.IS)]]
    if(length(IS) > 1) IS.is.var <- TRUE
  }
  # Report non-variables to user
  if(!T.is.var) {
    Tunits <- T.units()
    if(Tunits=="C") Tunits <- "\u00BAC"
    message('affinity: temperature is ', outvert(T, 'K'), ' ', Tunits)
  }
  if(!P.is.var) {
    if(identical(P,"Psat")) message("affinity: pressure is Psat")
    else message('affinity: pressure is ',outvert(P,'bar'),' ',P.units())
  }
  if(!IS.is.var & !identical(IS,0)) message('affinity: ionic strength is ',IS)
  # Default value for resolution
  res <- 256
  # Where we store the output
  what <- "A"
  vars <- character()
  vals <- list(NA)
  # This needs to have 1 as the third component because
  # energy() uses it to build an array with the given dimension
  lims <- list(c(NA, NA, 1))
  # Clean out non-variables
  if(any(arg.is.T) & !T.is.var) args <- args[names(args)!="T"]
  if(any(arg.is.P) & !P.is.var) args <- args[names(args)!="P"]
  if(any(arg.is.IS) & !IS.is.var) args <- args[names(args)!="IS"]
  # The property we're interested in
  if("what" %in% names(args)) {
    what <- args[[names(args)=="what"]]
    args <- args[names(args)!="what"]
  }
  # Assemble the variables
  if(length(args) > 0) {
    for(i in 1:length(args)) {
      nametxt <- names(args)[i]
      if(transect) lims.orig <- c(min(args[[i]]),max(args[[i]]))
      else lims.orig <- args[[i]][1:2]
      if(names(args)[i]=="pH") {
        names(args)[i] <- "H+"
        if(transect) args[[i]] <- -args[[i]]
        else args[[i]][1:2] <- -args[[i]][1:2]
      } 
      if(names(args)[i]=="pe") {
        names(args)[i] <- "e-"
        if(transect) args[[i]] <- -args[[i]]
        else args[[i]][1:2] <- -args[[i]][1:2]
      }
      if(length(args[[i]]) < 3 & !transect) args[[i]] <- c(args[[i]],res)
      vars[length(vars)+1] <- names(args)[i]
      if(transect) {
        vals[[length(vars)]] <- args[[i]]
        lims[[length(vars)]] <- c(lims.orig,length(vals[[i]]))
      } else {
        vals[[length(vars)]] <- seq(args[[i]][1],args[[i]][2],length.out=args[[i]][3])
        lims[[length(vars)]] <- args[[i]]
      }
      names(lims)[length(vars)] <- names(args)[i]
      # Say something about the identities, ranges, and units of the variables
      unittxt <- ""
      # Number of values
      if(transect) n <- length(args[[i]]) else n <- args[[i]][3]
      # Physical state
      ibasis <- match(nametxt, rownames(thermo$basis))
      if(isTRUE(as.logical(ibasis))) {
        if(thermo$basis$state[ibasis]=="gas") nametxt <- paste("log10(f_", nametxt, ")", sep="") 
        else nametxt <- paste("log10(a_", nametxt, ")", sep="") 
      } else {
        # Stop if the argument doesn't correspond to a basis species, T, P, or IS
        if(!nametxt %in% c("T", "P", "IS")) {
          if(! (nametxt=="pH" & 'H+' %in% rownames(thermo$basis) | nametxt %in% c("pe", "Eh") & 'e-' %in% rownames(thermo$basis))) {
            stop(nametxt, " is not one of T, P, or IS, and does not match any basis species")
          }
        }
      }
      # Temperature and pressure and Eh
      if(nametxt=="T") unittxt <- " K"
      if(nametxt=="P") unittxt <- " bar"
      if(nametxt=="Eh") unittxt <- " V"
      message("affinity: variable ", length(vars), " is ", nametxt, 
        " at ", n, " values from ", lims.orig[1], " to ", lims.orig[2], unittxt)
    }
  }
  args <- list(what=what,vars=vars,vals=vals,lims=lims,T=T,P=P,IS=IS,transect=transect)

  # Convert Eh to pe
  if("Eh" %in% args$vars) {
    # Get Eh into our dimensions
    Eh.args <- args
    # What variable is Eh
    Eh.var <- which(args$vars=="Eh")
    Eh.args$what <- args$vals[[Eh.var]]
    Eh.args$sout$out <- Eh.var
    Eh <- do.call("energy",Eh.args)
    # Get temperature into our dimensions
    T.args <- args  
    if("T" %in% args$vars) {
      T.var <- which(args$vars=="T")
      T.args$what <- args$vals[[T.var]]
    } else {
      T.var <- 1
      T.args$what <- T
    }
    T.args$sout$out <- T.var
    T <- do.call("energy",T.args)
    # Do the conversion on vectors
    mydim <- dim(Eh)
    Eh <- as.vector(Eh)
    T <- as.vector(T)
    pe <- convert(Eh,"pe",T=T)
    dim(pe) <- mydim
    # Update the arguments list
    args$vars[Eh.var] <- "e-"
    args$vals[[Eh.var]] <- -pe
  }
  return(args)
}

A.ionization <- function(iprotein, vars, vals, T=298.15, P="Psat", pH=7, transect=FALSE) {
  # A function to build a list of values of A/2.303RT of protein ionization
  # that can be used by energy(); 20120527 jmd
  # Some of the variables might not affect the values (e.g. logfO2)
  # What are the variables that affect the values
  T <- convert(T, "C")
  if(!is.na(iT <- match("T", vars))) T <- convert(vals[[iT]], "C")
  if(!is.na(iP <- match("P", vars))) P <- vals[[iP]]
  if(!is.na(iHplus <- match("H+", vars))) pH <- -vals[[iHplus]]
  # Is it a transect (single points) or a grid?
  if(transect) TPpH <- list(T=T, P=P, pH=pH)
  else {
    # Make a grid of all combinations
    # Put T, P, pH in same order as they show up vars
    egargs <- list(T=T, P=P, pH=pH)
    # order(c(NA, NA, NA)) might segfault in some versions of R (seen in R 2.15.3 on Linux)
    if(is.na(iT) & is.na(iP) & is.na(iHplus)) TPpHorder <- c(1, 2, 3)
    else TPpHorder <- order(c(iT, iP, iHplus))
    egargs <- c(egargs[TPpHorder], list(stringsAsFactors=FALSE))
    TPpH <- do.call(expand.grid, egargs)
    # Figure out the dimensions of T-P-pH (making sure to drop any that aren't in vars)
    TPpHdim <- numeric(3)
    TPpHdim[iT] <- length(T)
    TPpHdim[iP] <- length(P)
    TPpHdim[iHplus] <- length(pH)
    TPpHdim <- TPpHdim[!TPpHdim==0]
    # Figure out the dimensions of the other vars
    othervars <- vars[!vars %in% c("T", "P", "H+")]
    iother <- match(othervars, vars)
    otherdim <- sapply(vals, length)[iother]
    # If Eh was given to energy.args, values of pe were calculated in all dimensions
    # Figure out the original length of the Eh variable
    ieminus <- match("e-", vars[iother])
    if(!is.na(ieminus)) {
      otherdim[ieminus] <- NA
      edim <- dim(vals[[iother[ieminus]]])
      # Loop TPpHdim and otherdim, taking out each one from edim
      for(dim in c(TPpHdim, otherdim)) {
        id <- match(dim, edim)
        edim[id] <- NA
      }
      otherdim[ieminus] <- edim[!is.na(edim)]
    }
    # The permutation vector
    if(length(iother) > 0) allvars <- c(vars[-iother], vars[iother])
    else allvars <- vars
    perm <- match(vars, allvars)
  }
  # Initialize output list
  out <- vector("list", length(iprotein))
  # Get aa from iprotein
  aa <- pinfo(iprotein)
  # Calculate the values of A/2.303RT as a function of T-P-pH
  A <- ionize.aa(aa=aa, property="A", T=TPpH$T, P=TPpH$P, pH=TPpH$pH)
  if(transect) {
    # Just turn the values into a list
    for(i in 1:length(iprotein)) out[[i]] <- A[, i]
  } else for(i in 1:length(iprotein)) {
    # What are the values of ionization affinity for this protein
    thisA <- A[, i]
    # Apply the dimensions of T-P-pH
    tpphdim <- TPpHdim
    if(length(tpphdim)==0) tpphdim <- 1
    thisA <- array(thisA, tpphdim)
    # Grow into the dimensions of all vars
    alldim <- c(TPpHdim, otherdim)
    if(length(alldim)==0) alldim <- 1
    thisA <- array(thisA, alldim)
    # Now permute the array to put dimensions in same order as the variables
    thisA <- aperm(thisA, perm)
    # Store in output list
    out[[i]] <- thisA
  }
  return(out)
}